Publish Modilify Mk2 step 4400

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,7 @@
+.DS_Store
+__pycache__/
+*.py[cod]
+.venv/
+.cache/
+.env
+/.upload-fast.sh

LICENSE

@@ -0,0 +1,167 @@
+Modilify Non-Commercial Model License 1.0
+
+Copyright 2026 Modilify
+
+1. Definitions
+
+"Model" means the weights, configuration, inference code, tokenizer, processor,
+architecture definitions, and documentation distributed under or with this license.
+"Derivative Model" means any modified, fine-tuned, distilled, merged, quantized,
+extended, or otherwise adapted version of the Model.
+
+"Output" means any text, image, audio, video, 3D asset, latent state, synthetic
+data, or other content generated by executing, querying, or running the Model
+or a Derivative Model.
+
+"Non-Commercial Use" means any use, exercise of rights, or deployment conducted
+strictly for personal learning, educational purposes, academic research,
+non-profit scientific investigation, or non-commercial testing and evaluation,
+which does not directly or indirectly generate revenue, commercial advantage,
+or monetary compensation.
+
+"You" means the individual or legal entity exercising permissions under this
+license.
+
+"High-Risk Use" means a use that can materially affect a person's safety,
+liberty, access to essential services, employment, housing, credit, education,
+legal rights, or medical care, or that controls critical infrastructure,
+weapons, or large-scale biometric surveillance.
+
+2. Copyright Grant
+
+Subject to the terms and conditions of this license, Modilify grants You a
+worldwide, perpetual, non-exclusive, royalty-free, non-transferable, revocable
+(as provided herein) copyright license to use, reproduce, prepare derivative
+works of, publicly display, and distribute the Model and Derivative Models
+solely for Non-Commercial Use.
+
+3. Commercial Restrictions, Outputs, and Distillation
+
+a. Prohibition on Commercial Use: You shall not use, deploy, host, serve,
+   reproduce, sublicense, or distribute the Model or any Derivative Model for
+   any commercial purpose, including but not limited to:
+   i.   operating the Model or Derivative Models as a paid or revenue-generating
+        service (including API or cloud-hosted services);
+   ii.  incorporating the Model or Derivative Models into commercial products,
+        enterprise workflows, or profit-seeking internal operations; or
+   iii. selling, leasing, or monetizing the Model or Derivative Models in any
+        manner.
+
+b. Output Restrictions: You shall not use any Output generated by the Model or
+   Derivative Models for direct commercial monetization or commercial products.
+
+c. Distillation and Competitive Training Prohibition: You shall not use the Model,
+   Derivative Models, or any Outputs generated by them (including synthetic data,
+   distillation targets, or reasoning traces) to train, fine-tune, distill,
+   align, improve, or evaluate any machine learning model that is intended for
+   commercial use or that competes with the Model.
+
+d. Commercial Inquiries: Any commercial use of the Model or its derivatives
+   requires a separate written commercial agreement from Modilify. Inquiries
+   regarding commercial licensing may be submitted through the official project
+   repository or Modilify's public channels.
+
+4. Patent Grant
+
+Each contributor grants You a worldwide, perpetual, non-exclusive, royalty-free,
+irrevocable patent license, except as stated in this section, to make, have made,
+use, and otherwise transfer the Model solely for Non-Commercial Use, where the
+license applies only to those patent claims licensable by that contributor that
+are necessarily infringed by that contributor's contribution alone or in
+combination with the Model. If You institute patent litigation alleging that the
+Model or a contribution constitutes patent infringement, patent licenses granted
+to You under this license terminate as of the filing date.
+
+5. Conditions on Redistribution and Share-Alike
+
+If You distribute the Model or a Derivative Model, You must:
+
+a. provide recipients a copy of this license;
+b. retain all copyright, patent, trademark, and attribution notices;
+c. state clearly that You modified the Model and document material modifications;
+d. Share-Alike Requirement: license and distribute any Derivative Model solely
+   under this exact Modilify Non-Commercial Model License 1.0, without imposing
+   any additional restrictions or granting any broader rights than those provided
+   herein; and
+e. publish with the distributed model a reasonably accessible impact statement
+   (such as within a public model card or repository documentation) describing
+   intended uses, material limitations, evaluation scope, known safety risks,
+   and risk mitigations for the Derivative Model.
+
+The impact statement may be maintained in a public model card or equivalent
+document. You are not required to submit it separately to Modilify.
+
+6. Responsible Use and High-Risk Uses
+
+You must not use the Model or a Derivative Model:
+
+a. to develop, operate, or materially facilitate weapons, autonomous targeting,
+   or systems intended to cause physical harm;
+b. for unlawful mass surveillance, biometric identification without lawful
+   authority and appropriate safeguards, or social scoring that determines
+   access to rights or essential services;
+c. to exploit children or vulnerable persons, facilitate human trafficking, or
+   generate non-consensual intimate content;
+d. to impersonate a person or deceptively represent machine output as an
+   authentic human communication where the deception is reasonably likely to
+   cause material harm; or
+e. to make a final decision in a High-Risk Use without meaningful qualified
+   human review, proportionate testing, monitoring, appeal or correction paths,
+   and compliance with applicable law.
+
+Before deploying the Model in a High-Risk Use, You must perform safety and impact
+due diligence proportionate to foreseeable harm. At minimum, document the use
+context, evaluate relevant failure modes and affected groups, apply reasonable
+technical and organizational safeguards, monitor material incidents, and update
+or suspend the deployment when its residual risk is not reasonable. Research,
+testing, auditing, and defensive safety work are permitted when conducted with
+appropriate safeguards.
+
+7. Trademarks
+
+This license does not grant permission to use the trade names, trademarks,
+service marks, or product names of Modilify or any contributor, except as needed
+for reasonable and customary attribution or to describe the origin of the Model.
+
+8. Third-Party Components
+
+If the Model incorporates or links to any third-party components or libraries,
+those components remain subject to their respective applicable licenses and
+terms.
+
+9. Termination and Reinstatement
+
+Your rights terminate automatically if You materially violate this license and
+do not cure the violation within 30 days after becoming aware of it. Rights are
+reinstated upon timely cure unless a rights holder provides written notice of a
+substantially similar repeated violation. Sections intended by their nature to
+survive termination remain effective.
+
+10. Disclaimer of Warranty
+
+THE MODEL IS PROVIDED "AS IS," WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING WARRANTIES OF TITLE, NON-INFRINGEMENT,
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, ACCURACY, OR SAFETY. YOU ARE
+SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING
+THE MODEL AND ASSUME ALL RISKS ASSOCIATED WITH YOUR EXERCISE OF PERMISSIONS.
+
+11. Limitation of Liability
+
+TO THE MAXIMUM EXTENT PERMITTED BY LAW, NO COPYRIGHT HOLDER OR CONTRIBUTOR SHALL
+BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+CONSEQUENTIAL DAMAGES ARISING FROM THIS LICENSE OR THE USE OR INABILITY TO USE
+THE MODEL, HOWEVER CAUSED AND UNDER ANY THEORY OF LIABILITY, EVEN IF ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+12. Governing Law
+
+To the maximum extent permitted by applicable law, this license is governed by
+and construed in accordance with applicable laws, excluding conflict-of-law
+rules.
+
+13. Entire License; Severability
+
+This document states the complete Modilify license for the Model. If any
+provision is held to be invalid or unenforceable, that provision will be limited
+or eliminated to the minimum extent necessary, and the remaining provisions
+of this license will remain in full force and effect.

LICENSE.apache-2.0

@@ -0,0 +1,204 @@
+Copyright 2018- The Hugging Face team. All rights reserved.
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+

NOTICE.md

@@ -0,0 +1,7 @@
+# Attribution
+
+Modilify Mk2 uses the text backbone of Google DeepMind's [DiffusionGemma 26B-A4B-it](https://huggingface.co/google/diffusiongemma-26B-A4B-it), whose model card identifies its license as Apache 2.0. The vision tower is not included in this text-only release.
+
+The Modilify architecture, rolling-canvas generation, trajectory memory, and trained dense/expert adapters modify the upstream model. This release is the saved step 4400 checkpoint, converted to a self-contained Transformers model. Modilify's contributions use the [Modilify Non-Commercial Model License 1.0](LICENSE). This does not replace upstream rights and terms.
+
+The code integrates with Hugging Face Transformers (Copyright 2018–, The Hugging Face team; Apache 2.0) and uses PyTorch and Safetensors as installed dependencies. Those libraries retain their own licenses. The Apache license text is reproduced in [LICENSE.apache-2.0](LICENSE.apache-2.0).

README.md

@@ -0,0 +1,114 @@
+---
+license: other
+license_name: modilify-non-commercial-model-license-1.0
+license_link: LICENSE
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+  - diffusion
+  - mixture-of-experts
+  - custom-code
+  - modilify-mk2
+---
+
+# Modilify Mk2
+
+Modilify Mk2 is a text diffusion model with learned trajectory memory. Built on the DiffusionGemma 26B-A4B text backbone, it develops a response through repeated revisions of a rolling text canvas. Its memory follows both the current revision process and the content already committed to the response.
+
+This repository provides the complete model weights trained on 78M tokens, tokenizer, and a PyTorch / Transformers runtime.
+
+## Architecture
+
+Modilify Mk2 works on a 256-token canvas. During each denoising pass, it revises candidate tokens together and uses their confidence and uncertainty to select a prefix to commit. The canvas then advances to continue the response.
+
+Two timescales of GDN2 memory connect these passes. Working memory tracks changes within the canvas, with states for individual positions and the canvas as a whole. Persistent memory incorporates experience from committed tokens and carries it into later parts of the response. The model learns these memory components alongside dense and expert LoRA adapters.
+
+Training supervises response tokens across the valid canvas, with additional objectives for confidence calibration, prefix readiness, and response termination.
+
+| Property | Value |
+| --- | --- |
+| Backbone | DiffusionGemma 26B-A4B, text only |
+| Adaptation | Dense and expert LoRA |
+| Memory | Dual-timescale GDN2 trajectory memory |
+| Canvas size | 256 tokens |
+| Training step | 4400 |
+| Adaptation tokens | 78,332,410 (prompt + response, excluding padding) |
+| Weight precision | BF16, with selected memory parameters in FP32 |
+| Weight files | 16 Safetensors shards, approximately 48.23 GiB |
+| Framework | PyTorch / Hugging Face Transformers |
+
+## Quickstart
+
+From the downloaded model directory, use Python 3.12 and install the dependencies:
+
+```bash
+python -m pip install -r requirements.txt
+python inference.py --prompt "Why is the sky blue?" --stream
+```
+
+The CLI selects CUDA, Apple Silicon MPS, or CPU according to availability. Choose a device explicitly with `--device cuda`, `--device mps`, or `--device cpu`. In addition to the weights, allow memory for the prompt cache, trajectory states, and generation workspace. Local inference has been exercised on CPU and MPS.
+
+## Use with Transformers
+
+The model loads through `AutoModelForCausalLM`. Set `model_path` to the downloaded directory or `Modilify/Modilify-Mk2` on the Hub.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_path = "./"
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    trust_remote_code=True,
+    dtype="auto",
+    device_map=device,
+    attn_implementation="sdpa",
+).eval()
+
+messages = [{"role": "user", "content": "Explain how rainbows form."}]
+inputs = tokenizer.apply_chat_template(
+    messages,
+    enable_thinking=True,
+    add_generation_prompt=True,
+    return_tensors="pt",
+    return_dict=True,
+).to(model.device)
+
+with torch.inference_mode():
+    output = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        max_denoising_steps=256,
+        seed=42,
+    )
+
+response_ids = output.sequences[0, inputs.input_ids.shape[1]:]
+print(tokenizer.decode(response_ids, skip_special_tokens=True))
+```
+
+### Thinking and conversations
+
+The chat template supports system, user, and assistant turns. Pass the conversation as `messages` and set `enable_thinking=True` to enable reasoning before the answer. The CLI enables thinking by default; `--no-thinking` disables it.
+
+### Generation controls
+
+`max_new_tokens` sets the response length budget, while `max_denoising_steps` sets the number of canvas revisions. Both default to 256. Increase the budgets for longer responses; reaching either limit can leave a response unfinished. `output.stop_reason` reports how generation ended.
+
+Use `seed` for a single request or `seeds=[...]` for batch requests. For batches, set `tokenizer.padding_side = "left"` and pass the attention mask returned by the tokenizer. The returned `output.sequences` includes the input prompt, and `output.generated_lengths` gives each response's token count.
+
+Generation runs through the included diffusion implementation. The weights include the backbone, adapters, and memory components; `from_pretrained` loads them together. `dtype="auto"` preserves the model's mixed precision settings.
+
+## Use and limitations
+
+Modilify Mk2 is released for non-commercial research and experimentation with text diffusion, iterative response refinement, and learned generation memory. It accepts text input and produces text, including an optional reasoning channel.
+
+Broad capability and safety benchmark results are not yet available for this checkpoint. Generated responses can contain factual errors, repetition, or incomplete reasoning. Review outputs before using them for consequential decisions. The adaptation dataset is not included in this release.
+
+## License and credits
+
+The model builds on Google DeepMind's [DiffusionGemma](https://huggingface.co/google/diffusiongemma-26B-A4B-it) and integrates with Hugging Face Transformers. Attribution and upstream licenses are listed in [NOTICE.md](NOTICE.md).
+
+Modilify's contributions are available under the [Modilify Non-Commercial Model License 1.0](LICENSE). Commercial use requires a separate agreement.

adapters.py

@@ -0,0 +1,121 @@
+"""Unfused dense and expert adapters for the published schema25 checkpoint."""
+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
+    dropout: float
+    target_modules: tuple[str, ...]
+    expert_r: int = 8
+    expert_alpha: int = 8
+
+
+def unwrap_model(model: nn.Module) -> nn.Module:
+    current = model
+    while True:
+        if hasattr(current, "module"):
+            current = current.module
+            continue
+        if hasattr(current, "_orig_mod"):
+            current = current._orig_mod
+            continue
+        return current
+
+
+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.dropout = nn.Dropout(config.dropout) if config.dropout else nn.Identity()
+        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:
+        dropped = self.dropout(inputs)
+        if inputs.device.type == "mps" and torch.is_grad_enabled() and dropped.requires_grad:
+            update = torch.nn.functional.linear(
+                torch.nn.functional.linear(dropped.float(), self.lora_a.float()),
+                self.lora_b.float(),
+            ).to(dtype=inputs.dtype)
+        else:
+            update = torch.nn.functional.linear(
+                torch.nn.functional.linear(dropped, 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) or isinstance(module, LoRALinear):
+            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,
+            dropout=config.dropout,
+        )
+    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."""
+    raw = unwrap_model(model)
+    encoder = raw.model.encoder.language_model
+    decoder = raw.model.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))

chat_template.jinja

@@ -0,0 +1,387 @@
+{#
+  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.
+#}
+{%- macro format_parameters(properties, required, filter_keys=false) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if not filter_keys or key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'OBJECT' -%}
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                    }
+                {%- elif value is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
+                    }
+                {%- endif -%}
+                {%- if value['required'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is none -%}
+        {{- 'null' -}}
+    {%- elif argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{#- ===== SETUP ===== -#}
+{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}
+{%- set loop_messages = messages -%}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}
+    {{- '<|turn>system\n' -}}
+    {#- Inject Thinking token at the very top of the FIRST system turn -#}
+    {%- if enable_thinking -%}
+        {{- '<|think|>\n' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+    {%- if messages and messages[0]['role'] in ['system', 'developer'] -%}
+        {%- if messages[0]['content'] is string -%}
+            {{- messages[0]['content'] | trim -}}
+        {%- elif messages[0]['content'] is sequence -%}
+            {%- for item in messages[0]['content'] -%}
+                {{- item['text'] | trim + ' '-}}
+            {%- endfor -%}
+        {%- endif -%}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+    {%- if message['role'] != 'tool' -%}
+    {%- set ns.prev_message_type = None -%}
+    {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+    {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#}
+    {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}
+    {%- if not continue_same_model_turn -%}
+        {{- '<|turn>' + role + '\n' }}
+        
+    {%- endif -%}
+
+    {#- Render reasoning/reasoning_content as thinking channel -#}
+    {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+    {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or (preserve_thinking and message.get('tool_calls')) -%}
+    {%- if thinking_text and thinking_gate -%}
+        {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
+    {%- endif -%}
+
+            {%- if message.get('tool_calls') -%}
+                {%- for tool_call in message.get('tool_calls') -%}
+                    {%- set function = tool_call['function'] -%}
+                    {{- '<|tool_call>call:' + function['name'] + '{' -}}
+                    {%- if function['arguments'] is mapping -%}
+                        {%- set ns_args = namespace(found_first=false) -%}
+                        {%- for key, value in function['arguments'] | dictsort -%}
+                            {%- if ns_args.found_first %},{% endif -%}
+                            {%- set ns_args.found_first = true -%}
+                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+                        {%- endfor -%}
+                    {%- elif function['arguments'] is none -%}
+                    {%- else -%}
+                      {{- raise_exception(
+                          "chat_template: tool_calls[].function.arguments must be a "
+                          "JSON object (mapping), not a string. Deserialize arguments "
+                          "before passing to the template."
+                      ) -}}
+                    {%- endif -%}
+                    {{- '}<tool_call|>' -}}
+                {%- endfor -%}
+                {%- set ns.prev_message_type = 'tool_call' -%}
+            {%- endif -%}
+
+            {%- set ns_tr_out = namespace(flag=false) -%}
+            {%- if message.get('tool_responses') -%}
+                {#- 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 -%}
+                    {%- set ns.prev_message_type = 'tool_response' -%}
+                {%- endfor -%}
+            {%- elif message.get('tool_calls') -%}
+                {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+                {%- set ns_tool_scan = namespace(stopped=false) -%}
+                {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+                    {%- if ns_tool_scan.stopped -%}
+                    {%- elif loop_messages[k]['role'] != 'tool' -%}
+                        {%- set ns_tool_scan.stopped = true -%}
+                    {%- else -%}
+                        {%- set follow = loop_messages[k] -%}
+                        {#- Resolve tool_call_id to function name -#}
+                        {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}
+                        {%- for tc in message.get('tool_calls') -%}
+                            {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                                {%- set ns_tname.name = tc['function']['name'] -%}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {#- Handle content as string or content-parts array -#}
+                        {%- set tool_body = follow.get('content') -%}
+                        {%- if tool_body is string -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- elif tool_body is sequence and tool_body is not string -%}
+                            {%- set ns_txt = namespace(s='') -%}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'text' -%}
+                                    {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                                {%- endif -%}
+                            {%- endfor -%}
+                            {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') in ['image', 'image_url'] -%}
+                                    {{- '<|image|>' -}}
+                                {%- elif part.get('type') in ['audio', 'input_audio'] -%}
+                                    {{- '<|audio|>' -}}
+                                {%- elif part.get('type') == 'video' -%}
+                                    {{- '<|video|>' -}}
+                                {%- endif -%}
+                            {%- endfor -%}
+                        {%- else -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- endif -%}
+                        {%- set ns_tr_out.flag = true -%}
+                        {%- set ns.prev_message_type = 'tool_response' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+            {%- set captured_content -%}
+            {%- if message.get('content') is string -%}
+                {%- if role == 'model' -%}
+                    {{- strip_thinking(message['content']) -}}
+                {%- else -%}
+                    {{- message['content'] | trim -}}
+                {%- endif -%}
+            {%- elif message.get('content') is sequence -%}
+                {%- for item in message['content'] -%}
+                    {%- if item.get('type') == 'text' -%}
+                        {%- if role == 'model' -%}
+                            {{- strip_thinking(item['text']) -}}
+                        {%- else -%}
+                            {{- item['text'] | trim -}}
+                        {%- endif -%}
+                    {%- elif item.get('type') in ['image', 'image_url'] -%}
+                        {{- '<|image|>' -}}
+                    {%- elif item.get('type') in ['audio', 'input_audio'] -%}
+                        {{- '<|audio|>' -}}
+                    {%- elif item.get('type') == 'video' -%}
+                        {{- '<|video|>' -}}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+            {%- endset -%}
+
+            {{- captured_content -}}
+            {%- set has_content = captured_content | trim | length > 0 -%}
+
+        {#- Forward-scan: find next non-tool message role for continuation detection -#}
+        {%- set next_nt = namespace(role=None, found=false) -%}
+        {%- for j in range(loop.index0 + 1, loop_messages | length) -%}
+            {%- if not next_nt.found -%}
+                {%- if loop_messages[j]['role'] != 'tool' -%}
+                    {%- set next_nt.role = loop_messages[j]['role'] -%}
+                    {%- set next_nt.found = true -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+
+        {%- set continues_into_next = (
+            role == 'model'
+            and next_nt.role == 'assistant'
+            and (not message.get('tool_calls') or ns_tr_out.flag)
+        ) -%}
+        
+        {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+            {{- '<|tool_response>' -}}
+        {%- elif continues_into_next -%}
+        {%- elif not (ns_tr_out.flag and not has_content and not next_nt.found) -%}
+            {{- '<turn|>\n' -}}
+        {%- endif -%}
+    {%- endif -%}
+
+    {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}
+    {%- set ns.prev_non_tool_role = message['role'] -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+        {{- '<|turn>model\n' -}}
+    {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%}
+        {{- '<|channel>thought\n' -}}
+    {%- endif -%}
+{%- endif -%}

commit_policy.py

@@ -0,0 +1,370 @@
+"""Shared confidence-and-entropy commit policy for training and inference."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import math
+import torch
+
+from .latent_deliberation import (
+    advance_trajectory_clocks,
+    should_force_trajectory_jump,
+)
+
+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,
+    *,
+    eps: float = FUSED_EPS,
+    entropy_weight: float = 1.0,
+    confidence_power: float = 2.0,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    target_confidence: float | None = None,
+    failure_budget: float = 0.2,
+    **_kwargs: object,
+) -> torch.Tensor:
+    """Fuse proposal confidence with token entropy into effective commit confidence.
+
+    Effective confidence is defined using an excess-entropy sigmoid:
+
+        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) + bias - entropy_weight * excess)
+        fused = base_fused ** confidence_power
+
+    The helper defaults preserve legacy callers. Runtime entrypoints pass the
+    checkpoint's explicit policy parameters; new runs default to power one.
+    """
+
+    p = proposal_confidence.float().nan_to_num(0.5).clamp(min=eps, max=1.0 - eps)
+    H = token_entropy.float().nan_to_num(0.0).clamp(min=0.0)
+
+    # Binary entropy of p, using log1p for the (1-p) term.
+    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)
+
+    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)
+
+
+def sampled_commit_log_risk(
+    sampled_nll: torch.Tensor, token_entropy: torch.Tensor, *, entropy_weight: float,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    target_confidence: float | None = None,
+    failure_budget: float = 0.2,
+    **_kwargs: object,
+) -> torch.Tensor:
+    """Stable -log(base fused confidence), without probability saturation.
+
+    The power-normalized proxy is shared across confidence-power ablations.
+    NLL comes directly from log Z minus the sampled score, not log(clamp(p)).
+    """
+    raw_nll = sampled_nll.float()
+    # FP32 reductions can produce a tiny negative log Z - selected score.
+    # Correct the forward domain without clipping the sampled-score gradient.
+    nll = raw_nll + (raw_nll.detach().clamp_min(0.0) - raw_nll.detach())
+    p = torch.exp(-nll)
+    complement = -torch.expm1(-nll)
+    h2 = p * nll - complement * torch.log(complement.clamp_min(1e-12))
+    excess = (token_entropy.float() - 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(complement / 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:
+        excess = torch.minimum(excess, complement * torch.log(k_eff))
+    bias = commit_target_confidence_bias(target_confidence, failure_budget=failure_budget)
+    u = float(entropy_weight) * excess - bias
+    small = torch.log1p(complement * torch.expm1(u.clamp_max(20.0)))
+    large = u + torch.log((complement + p * torch.exp(-u)).clamp_min(1e-12))
+    return nll + torch.where(u <= 20.0, small, large)
+
+
+def fused_commit_failure_rate(
+    proposal_confidence: torch.Tensor,
+    token_entropy: torch.Tensor,
+    **kwargs: object,
+) -> torch.Tensor:
+    """Return ``1 - effective_confidence`` from the shared fusion helper."""
+
+    return 1.0 - fused_commit_confidence(
+        proposal_confidence, token_entropy, **kwargs
+    )
+
+
+@dataclass(frozen=True)
+class CommitPolicyDecision:
+    """One shared training/inference transition from proposal to commit."""
+
+    normal_lengths: torch.LongTensor
+    commit_lengths: torch.LongTensor
+    commit_token_ids: torch.LongTensor
+    jump_rows: torch.BoolTensor
+    ponder_steps: torch.IntTensor
+    stagnation_steps: torch.IntTensor
+
+
+def prefix_failure_commit_lengths(
+    failure_rate: torch.Tensor,
+    *,
+    failure_budget: float,
+    valid_mask: torch.BoolTensor | None = None,
+) -> torch.LongTensor:
+    """Return the longest valid prefix satisfying ``cumsum(failure_rate) < budget``."""
+
+    if failure_rate.ndim != 2:
+        raise ValueError("Failure rate must have shape [batch, canvas].")
+    if failure_budget <= 0:
+        raise ValueError("Commit failure budget must be positive.")
+    if valid_mask is None:
+        valid_mask = torch.ones_like(failure_rate, dtype=torch.bool)
+    if valid_mask.shape != failure_rate.shape:
+        raise ValueError("Commit validity mask must match failure rate.")
+
+    risk = failure_rate.float().clamp(0.0, 1.0) * valid_mask.to(torch.float32)
+    cumulative_risk = risk.cumsum(dim=-1)
+    contiguous_valid = valid_mask.long().cumprod(dim=-1).bool()
+    allowed = cumulative_risk.lt(float(failure_budget)) & contiguous_valid
+    return allowed.long().cumprod(dim=-1).sum(dim=-1)
+
+
+def first_committed_token_lengths(
+    proposal: torch.LongTensor,
+    commit_lengths: torch.LongTensor,
+    token_id: int | Sequence[int],
+    *,
+    positions: torch.LongTensor | None = None,
+) -> torch.LongTensor:
+    """Clip each committed prefix immediately after its first matching stop token."""
+
+    if proposal.ndim != 2 or commit_lengths.shape != proposal.shape[:1]:
+        raise ValueError("Proposal and commit lengths must share a batch dimension.")
+    if positions is None:
+        positions = torch.arange(proposal.shape[1], device=proposal.device).unsqueeze(0)
+    elif positions.shape != (1, proposal.shape[1]):
+        raise ValueError("Commit positions must have shape [1, canvas].")
+    committed = positions.lt(commit_lengths[:, None])
+    stop_token_ids = (
+        (int(token_id),)
+        if isinstance(token_id, int)
+        else tuple(dict.fromkeys(int(value) for value in token_id))
+    )
+    if not stop_token_ids:
+        raise ValueError("At least one stop token ID is required.")
+    matches = proposal.eq(stop_token_ids[0])
+    for value in stop_token_ids[1:]:
+        matches |= proposal.eq(value)
+    matches &= committed
+    sentinel = torch.full_like(positions, proposal.shape[1])
+    first = torch.where(matches, positions, sentinel).min(dim=-1).values
+    clipped = torch.where(first.lt(proposal.shape[1]), first + 1, commit_lengths)
+    return torch.minimum(clipped, commit_lengths)
+
+
+def bounded_prefix_failure_commit_lengths(
+    committed_token_ids: torch.LongTensor,
+    failure_rate: torch.Tensor,
+    *,
+    failure_budget: float,
+    remaining_lengths: torch.LongTensor,
+    stop_token_id: int | Sequence[int],
+    valid_mask: torch.BoolTensor | None = None,
+    positions: torch.LongTensor | None = None,
+) -> torch.LongTensor:
+    """Apply length and stop-token bounds to the shared failure-rate policy."""
+
+    if committed_token_ids.shape != failure_rate.shape:
+        raise ValueError("Committed token IDs and failure rate must share [batch, canvas].")
+    if remaining_lengths.shape != committed_token_ids.shape[:1]:
+        raise ValueError("Remaining lengths must have shape [batch].")
+    commit_lengths = prefix_failure_commit_lengths(
+        failure_rate,
+        failure_budget=failure_budget,
+        valid_mask=valid_mask,
+    )
+    commit_lengths = torch.minimum(commit_lengths, remaining_lengths.clamp_min(0))
+    return first_committed_token_lengths(
+        committed_token_ids,
+        commit_lengths,
+        stop_token_id,
+        positions=positions,
+    )
+
+
+def select_commit_lengths(
+    sampled_token_ids: torch.LongTensor,
+    normal_failure_rate: torch.Tensor,
+    previous_failure_rate: torch.Tensor,
+    greedy_token_ids: torch.LongTensor,
+    jump_failure_rate: torch.Tensor,
+    *,
+    ponder_steps: torch.Tensor,
+    stagnation_steps: torch.Tensor,
+    active_rows: torch.BoolTensor,
+    remaining_lengths: torch.LongTensor,
+    failure_budget: float,
+    stop_token_id: int | Sequence[int],
+    stagnation_threshold: int,
+    min_progress: float,
+    max_ponder_steps: int | None = None,
+    valid_mask: torch.BoolTensor | None = None,
+) -> CommitPolicyDecision:
+    """Use normal sampled commits and a fixed-budget greedy JUMP.
+
+    Progress is measured from the signed change in fused failure rate over the
+    frontier region (the union of the previous and current commit prefixes plus
+    one blocking position), not from raw confidence/entropy deltas.
+    """
+
+    if not (
+        sampled_token_ids.shape
+        == normal_failure_rate.shape
+        == previous_failure_rate.shape
+        == greedy_token_ids.shape
+        == jump_failure_rate.shape
+    ):
+        raise ValueError("Sampled and greedy statistics must share [batch, canvas].")
+
+    canvas_length = normal_failure_rate.shape[1]
+    positions = torch.arange(canvas_length, device=normal_failure_rate.device)[None, :]
+    normal = bounded_prefix_failure_commit_lengths(
+        sampled_token_ids,
+        normal_failure_rate,
+        failure_budget=failure_budget,
+        remaining_lengths=remaining_lengths,
+        stop_token_id=stop_token_id,
+        valid_mask=valid_mask,
+        positions=positions,
+    )
+    previous_prefix_length = prefix_failure_commit_lengths(
+        previous_failure_rate,
+        failure_budget=failure_budget,
+        valid_mask=valid_mask,
+    )
+    frontier_length = torch.maximum(previous_prefix_length, normal) + 1
+    valid_lengths = (
+        valid_mask.long().sum(dim=-1)
+        if valid_mask is not None
+        else torch.full_like(frontier_length, canvas_length)
+    )
+    frontier_length = torch.minimum(frontier_length, valid_lengths)
+    progress_mask = positions < frontier_length[:, None]
+    if valid_mask is not None:
+        progress_mask &= valid_mask
+    progress_mask &= active_rows[:, None]
+    signed_improvement = (
+        previous_failure_rate.float() - normal_failure_rate.float()
+    )
+    weights = progress_mask.float()
+    progress = (
+        signed_improvement * weights
+    ).sum(dim=-1) / weights.sum(dim=-1).clamp_min(1.0)
+    next_ponder, next_stagnation = advance_trajectory_clocks(
+        ponder_steps,
+        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,
+        progress_scores=progress,
+        min_progress=min_progress,
+        stagnation_threshold=stagnation_threshold,
+        ponder_steps=next_ponder,
+        max_ponder_steps=max_ponder_steps,
+    )
+    jump_commit = bounded_prefix_failure_commit_lengths(
+        greedy_token_ids,
+        jump_failure_rate,
+        failure_budget=JUMP_FAILURE_BUDGET,
+        remaining_lengths=remaining_lengths,
+        stop_token_id=stop_token_id,
+        valid_mask=valid_mask,
+        positions=positions,
+    )
+    committed = torch.where(jump_rows, jump_commit, normal)
+    commit_token_ids = torch.where(
+        jump_rows[:, None],
+        greedy_token_ids,
+        sampled_token_ids,
+    )
+    committed = torch.where(active_rows, committed, 0)
+    committed_rows = committed.gt(0)
+    jump_rows &= committed_rows
+    next_ponder = torch.where(
+        committed_rows,
+        0,
+        next_ponder,
+    ).to(torch.int32)
+    next_stagnation = torch.where(
+        committed_rows,
+        0,
+        next_stagnation,
+    ).to(torch.int32)
+    return CommitPolicyDecision(
+        normal_lengths=normal,
+        commit_lengths=committed,
+        commit_token_ids=commit_token_ids,
+        jump_rows=jump_rows,
+        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",
+    "sampled_commit_log_risk",
+    "prefix_failure_commit_lengths",
+    "select_commit_lengths",
+]

config.json

@@ -0,0 +1,161 @@
+{
+  "architectures": [
+    "ModilifyMk2ForBlockDiffusion"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_modilify_mk2.ModilifyMk2Config",
+    "AutoModel": "modeling_modilify_mk2.ModilifyMk2ForBlockDiffusion",
+    "AutoModelForCausalLM": "modeling_modilify_mk2.ModilifyMk2ForBlockDiffusion"
+  },
+  "canvas_length": 256,
+  "channel_end_token_id": 101,
+  "checkpoint_step": 4400,
+  "commit_confidence_power": 1.0,
+  "commit_entropy_weight": 1.0,
+  "commit_failure_budget": 0.2,
+  "commit_gold_alpha": 0.4,
+  "commit_gold_weight": 0.2,
+  "commit_min_p": 0.05,
+  "commit_readiness_beta": 0.02,
+  "commit_readiness_budget_margin": 0.02,
+  "commit_readiness_failure_weight": 5.0,
+  "commit_readiness_loss_weight": 0.5,
+  "commit_readiness_objective": "frontier_prefix_budget_v2",
+  "commit_readiness_target_tokens": 16,
+  "commit_sequence_dim": 1024,
+  "commit_sequence_layers": 2,
+  "commit_target_confidence": 0.5,
+  "commit_top_k": 40,
+  "confidence_calibration_loss_weight": 0.01,
+  "dtype": "bfloat16",
+  "eos_token_id": 1,
+  "initializer_range": 0.02,
+  "kv_cache_bucket_size": 128,
+  "latent_dim": 2816,
+  "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_persistent_bus_unfreeze_steps": 0,
+  "latent_tape_probes": 4,
+  "latent_tape_scheme": "gdn2_spatial_probe_v1",
+  "latent_working_bus_unfreeze_steps": 0,
+  "latent_working_last_block_global": true,
+  "lora_config": {
+    "alpha": 16,
+    "dropout": 0.05,
+    "expert_alpha": 8,
+    "expert_r": 8,
+    "r": 16,
+    "target_modules": [
+      "q_proj",
+      "k_proj",
+      "v_proj",
+      "o_proj",
+      "gate_proj",
+      "up_proj",
+      "down_proj",
+      "proj"
+    ]
+  },
+  "memory_architecture": "compact_gdn2_v2",
+  "memory_scheme": "dual_timescale_gdn2_trajectory_memory",
+  "model_type": "modilify_mk2",
+  "persistent_memory_bus": true,
+  "persistent_memory_write": "commit_only_transformer",
+  "state_schema_version": 25,
+  "terminal_stop_loss_weight": 0.01,
+  "terminal_stop_target_probability": 0.9,
+  "terminal_token_ids": [
+    106,
+    50
+  ],
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 2,
+    "dtype": "bfloat16",
+    "eos_token_id": 1,
+    "final_logit_softcapping": 30.0,
+    "global_head_dim": 512,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 2816,
+    "initializer_range": 0.02,
+    "intermediate_size": 2112,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "model_type": "modilify_mk2_text",
+    "moe_intermediate_size": 704,
+    "num_attention_heads": 16,
+    "num_experts": 128,
+    "num_global_key_value_heads": 2,
+    "num_hidden_layers": 30,
+    "num_key_value_heads": 8,
+    "pad_token_id": 0,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.25,
+        "rope_theta": 1000000.0,
+        "rope_type": "proportional"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 1024,
+    "tie_word_embeddings": true,
+    "top_k_experts": 8,
+    "use_bidirectional_attention": null,
+    "vocab_size": 262144
+  },
+  "tie_word_embeddings": true,
+  "token_ce_normalization": "per_sample_exposure_v1",
+  "token_ce_supervision": "valid_canvas_v1",
+  "token_loss_weight": 1.0,
+  "training_bptt_steps": 16,
+  "training_prefix_cache": "incremental",
+  "training_scheme": "gold_prefix_shared_commit_0_256_committed_ce_calibration_causal_throughput_terminal_sft",
+  "transformers_version": "5.14.1",
+  "turn_end_token_id": 106,
+  "vocab_chunk_size": 32768,
+  "working_memory_bus": true,
+  "writer_slot_gate": "per_slot"
+}

configuration_modilify_mk2.py

@@ -0,0 +1,352 @@
+"""Strict text-only configuration for the schema25 GDN2 protocol."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from transformers.configuration_utils import PreTrainedConfig
+
+from transformers.models.diffusion_gemma import DiffusionGemmaTextConfig
+
+
+STATE_SCHEMA_VERSION = 25
+# Keep topology stable; CE scope and normalization have separate objective tags.
+TRAINING_SCHEME = (
+    "gold_prefix_shared_commit_0_256_committed_ce_calibration_"
+    "causal_throughput_terminal_sft"
+)
+MEMORY_SCHEME = "dual_timescale_gdn2_trajectory_memory"
+MEMORY_ARCHITECTURE = "compact_gdn2_v2"
+CONFIG_PROTOCOL_ERROR = "ModilifyMk2 requires schema25 compact_gdn2_v2; older memory topologies require a new run."
+HISTORY_VIEWS = 4
+COMMIT_SEQUENCE_LAYERS = 2
+DENOISE_TEMPERATURE = 0.8
+VOCAB_CHUNK_SIZE = 32_768
+COMMIT_READINESS_OBJECTIVE = "frontier_prefix_budget_v2"
+TOKEN_CE_SUPERVISION = "valid_canvas_v1"
+TOKEN_CE_NORMALIZATION = "per_sample_exposure_v1"
+
+
+def require_current_checkpoint_protocol(metadata: Mapping[str, Any]) -> None:
+    """Reject checkpoints that do not implement the GDN2 state topology."""
+
+    if (
+        metadata.get("state_schema_version") != STATE_SCHEMA_VERSION
+        or metadata.get("training_scheme") != TRAINING_SCHEME
+        or metadata.get("memory_scheme") != MEMORY_SCHEME
+        or metadata.get("memory_architecture") != MEMORY_ARCHITECTURE
+    ):
+        raise RuntimeError(
+            f"ModilifyMk2 checkpoint requires schema{STATE_SCHEMA_VERSION} {MEMORY_ARCHITECTURE}; "
+            "older memory topologies cannot be restored. Start from the base model."
+        )
+
+
+class ModilifyMk2TextConfig(DiffusionGemmaTextConfig):
+    model_type = "modilify_mk2_text"
+    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):
+    """Configuration with recurrent GDN2 trajectory memory."""
+
+    model_type = "modilify_mk2"
+    sub_configs = {"text_config": ModilifyMk2TextConfig}
+
+    def __init__(
+        self,
+        text_config: ModilifyMk2TextConfig | dict[str, Any] | None = None,
+        *,
+        canvas_length: int = 256,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = True,
+        state_schema_version: int = STATE_SCHEMA_VERSION,
+        training_scheme: str = TRAINING_SCHEME,
+        memory_scheme: str = MEMORY_SCHEME,
+        memory_architecture: str = MEMORY_ARCHITECTURE,
+        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_history_length: int = 16,
+        latent_tape_probes: int = 1,
+        latent_tape_scheme: str = "gdn2_spatial_probe_v1",
+        latent_history_views: int = HISTORY_VIEWS,
+        latent_history_kv_rank: int | None = None,
+        latent_working_last_block_global: bool = True,
+        working_memory_bus: bool = True,
+        persistent_memory_bus: bool = True,
+        persistent_memory_write: str = "commit_only_transformer",
+        commit_sequence_layers: int = COMMIT_SEQUENCE_LAYERS,
+        commit_sequence_dim: int | None = None,
+        writer_slot_gate: str = "per_slot",
+        latent_working_bus_unfreeze_steps: int = 0,
+        latent_persistent_bus_unfreeze_steps: int = 0,
+        training_bptt_steps: int = 16,
+        kv_cache_bucket_size: int = 128,
+        turn_end_token_id: int = 106,
+        terminal_token_ids: Sequence[int] | None = None,
+        channel_end_token_id: int = 101,
+        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,
+        commit_gold_alpha: float = 0.4,
+        commit_gold_weight: float = 1.0,
+        commit_readiness_failure_weight: float = 1.0,
+        token_loss_weight: float = 1.0,
+        token_ce_supervision: str = TOKEN_CE_SUPERVISION,
+        token_ce_normalization: str = TOKEN_CE_NORMALIZATION,
+        confidence_calibration_loss_weight: float = 0.1,
+        commit_readiness_objective: str = COMMIT_READINESS_OBJECTIVE,
+        commit_readiness_target_tokens: int = 16,
+        commit_readiness_loss_weight: float = 0.1,
+        commit_readiness_budget_margin: float = 0.02,
+        commit_readiness_beta: float = 0.02,
+        terminal_stop_loss_weight: float = 1.0,
+        terminal_stop_target_probability: float = 0.95,
+        **kwargs: Any,
+    ) -> None:
+        if any(key.startswith("commit_throughput_") for key in kwargs):
+            raise ValueError("Removed commit-throughput configuration fields.")
+        if token_ce_supervision != TOKEN_CE_SUPERVISION:
+            raise ValueError("Schema25 requires valid-canvas token CE.")
+        if state_schema_version != STATE_SCHEMA_VERSION:
+            raise RuntimeError(CONFIG_PROTOCOL_ERROR)
+        if training_scheme != TRAINING_SCHEME or memory_scheme != MEMORY_SCHEME:
+            raise RuntimeError(CONFIG_PROTOCOL_ERROR)
+        if memory_architecture != MEMORY_ARCHITECTURE:
+            raise RuntimeError("Schema25 requires compact_gdn2_v2 memory; start a new run.")
+        if text_config is None:
+            text_config = ModilifyMk2TextConfig()
+        elif isinstance(text_config, dict):
+            text_config = dict(text_config)
+            text_config.pop("model_type", None)
+            text_config["use_bidirectional_attention"] = None
+            text_config = ModilifyMk2TextConfig(**text_config)
+        elif not isinstance(text_config, ModilifyMk2TextConfig):
+            payload = text_config.to_dict()
+            payload.pop("model_type", None)
+            text_config = ModilifyMk2TextConfig(**payload)
+
+        self.text_config = text_config
+        self.canvas_length = canvas_length
+        self.initializer_range = initializer_range
+        self.state_schema_version = STATE_SCHEMA_VERSION
+        self.training_scheme = training_scheme
+        self.memory_scheme = memory_scheme
+        self.memory_architecture = memory_architecture
+        self.latent_dim = latent_dim
+        self.latent_ffn_dim = latent_ffn_dim
+        self.latent_memory_slots = latent_memory_slots
+        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_history_length = latent_history_length
+        self.latent_tape_probes = int(latent_tape_probes)
+        self.latent_tape_scheme = str(latent_tape_scheme)
+        self.latent_history_views = int(latent_history_views)
+        if latent_history_kv_rank is None:
+            rank = min(1024, latent_dim)
+            rank -= rank % max(latent_num_heads, 1)
+            if rank <= 0:
+                rank = latent_num_heads
+            self.latent_history_kv_rank = rank
+        else:
+            self.latent_history_kv_rank = 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.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.latent_working_bus_unfreeze_steps = int(latent_working_bus_unfreeze_steps)
+        self.latent_persistent_bus_unfreeze_steps = int(
+            latent_persistent_bus_unfreeze_steps
+        )
+        self.training_bptt_steps = training_bptt_steps
+        self.kv_cache_bucket_size = kv_cache_bucket_size
+        self.turn_end_token_id = turn_end_token_id
+        if terminal_token_ids is None:
+            self.terminal_token_ids = (int(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.commit_failure_budget = float(commit_failure_budget)
+        self.commit_top_k = int(commit_top_k) if commit_top_k is not None else None
+        self.commit_min_p = float(commit_min_p) if commit_min_p is not None else None
+        self.commit_target_confidence = float(commit_target_confidence) if commit_target_confidence is not None else None
+        self.commit_entropy_weight = float(commit_entropy_weight)
+        self.commit_confidence_power = float(commit_confidence_power)
+        self.commit_gold_alpha = float(commit_gold_alpha)
+        self.commit_gold_weight = float(commit_gold_weight)
+        self.commit_readiness_failure_weight = float(commit_readiness_failure_weight)
+        self.token_loss_weight = token_loss_weight
+        self.token_ce_supervision = str(token_ce_supervision)
+        self.token_ce_normalization = str(token_ce_normalization)
+        self.confidence_calibration_loss_weight = confidence_calibration_loss_weight
+        self.commit_readiness_objective = str(commit_readiness_objective)
+        self.commit_readiness_target_tokens = int(commit_readiness_target_tokens)
+        self.commit_readiness_loss_weight = float(commit_readiness_loss_weight)
+        self.commit_readiness_budget_margin = float(commit_readiness_budget_margin)
+        self.commit_readiness_beta = float(commit_readiness_beta)
+        self.terminal_stop_loss_weight = terminal_stop_loss_weight
+        self.terminal_stop_target_probability = terminal_stop_target_probability
+        self.vocab_chunk_size = VOCAB_CHUNK_SIZE
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+        self._validate()
+
+    def _validate(self) -> None:
+        positive = (
+            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.training_bptt_steps, self.kv_cache_bucket_size,
+        )
+        if any(value <= 0 for value in positive):
+            raise ValueError("All schema25 dimensions and intervals must be positive.")
+        if self.canvas_length != 256:
+            raise ValueError("Schema25 requires a 256-token canvas.")
+        if self.latent_tape_scheme != "gdn2_spatial_probe_v1":
+            raise ValueError("Schema25 requires GDN2 spatial probes.")
+        if self.latent_tape_probes > self.canvas_length:
+            raise ValueError("Spatial probes cannot exceed canvas positions.")
+        if self.commit_sequence_dim % self.latent_num_heads:
+            raise ValueError("`commit_sequence_dim` must be divisible by `latent_num_heads`.")
+        if self.latent_working_bus_unfreeze_steps < 0:
+            raise ValueError("`latent_working_bus_unfreeze_steps` must be non-negative.")
+        if self.latent_persistent_bus_unfreeze_steps < 0:
+            raise ValueError(
+                "`latent_persistent_bus_unfreeze_steps` must be non-negative."
+            )
+        if self.latent_persistent_bus_unfreeze_steps < self.latent_working_bus_unfreeze_steps:
+            raise ValueError(
+                "Persistent bus must unfreeze no earlier than the working bus."
+            )
+        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.eos_token_id, int) or self.eos_token_id < 0:
+            raise ValueError("ModilifyMk2 requires one non-negative integer EOS token ID.")
+        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 self.commit_readiness_objective != COMMIT_READINESS_OBJECTIVE:
+            raise ValueError(
+                "Unsupported commit-readiness objective: "
+                f"{self.commit_readiness_objective!r}."
+            )
+        import math
+        if not all(math.isfinite(v) for v in (
+            self.commit_failure_budget, self.commit_entropy_weight,
+            self.commit_confidence_power, self.commit_gold_alpha, self.commit_gold_weight,
+            self.commit_readiness_failure_weight, self.commit_readiness_loss_weight,
+        )):
+            raise ValueError("Commit policy parameters must be finite.")
+        if self.commit_failure_budget <= 0 or self.commit_entropy_weight < 0:
+            raise ValueError("Commit budget must be positive and entropy weight non-negative.")
+        if self.commit_top_k is not None and self.commit_top_k <= 0:
+            raise ValueError("`commit_top_k` must be a positive integer.")
+        if self.commit_min_p is not None and not 0.0 < self.commit_min_p < 1.0:
+            raise ValueError("`commit_min_p` must be in (0, 1).")
+        if self.commit_target_confidence is not None and not 0.0 < self.commit_target_confidence < 1.0:
+            raise ValueError("`commit_target_confidence` must be in (0, 1).")
+        if self.commit_confidence_power <= 0 or not 0 < self.commit_gold_alpha < 1:
+            raise ValueError("Commit power must be positive and gold alpha in (0, 1).")
+        if self.commit_gold_weight <= 0:
+            raise ValueError("Commit gold weight must be positive.")
+        if self.commit_readiness_failure_weight < 0:
+            raise ValueError("Commit readiness failure weight must be non-negative.")
+        if self.commit_readiness_target_tokens <= 0:
+            raise ValueError("`commit_readiness_target_tokens` must be positive.")
+        if self.commit_readiness_loss_weight < 0:
+            raise ValueError("`commit_readiness_loss_weight` must be non-negative.")
+        if not 0.0 <= self.commit_readiness_budget_margin < self.commit_failure_budget:
+            raise ValueError(
+                "`commit_readiness_budget_margin` must be in [0, commit_failure_budget)."
+            )
+        if self.commit_readiness_beta <= 0:
+            raise ValueError("`commit_readiness_beta` must be positive.")
+        if self.terminal_stop_loss_weight < 0:
+            raise ValueError("`terminal_stop_loss_weight` must be non-negative.")
+        if not 0.0 < self.terminal_stop_target_probability < 1.0:
+            raise ValueError("`terminal_stop_target_probability` must be in (0, 1).")
+        loss_weights = (
+            self.token_loss_weight,
+            self.confidence_calibration_loss_weight,
+            self.commit_readiness_loss_weight,
+            self.terminal_stop_loss_weight,
+        )
+        if any(not math.isfinite(weight) or weight < 0 for weight in loss_weights):
+            raise ValueError("Schema25 fixed loss weights must be finite and non-negative.")
+        if self.token_ce_normalization != TOKEN_CE_NORMALIZATION:
+            raise ValueError("Unsupported token CE normalization.")
+
+    @classmethod
+    def from_dict(cls, config_dict: dict[str, Any], **kwargs: Any) -> "ModilifyMk2Config":
+        return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)
+        payload = dict(config_dict)
+        if (payload.get("state_schema_version") != STATE_SCHEMA_VERSION
+                or payload.get("memory_architecture") != MEMORY_ARCHITECTURE):
+            raise RuntimeError(CONFIG_PROTOCOL_ERROR)
+        payload.pop("model_type", None)
+        payload.pop("architectures", None)
+        for key in tuple(kwargs):
+            if key in payload:
+                payload[key] = kwargs.pop(key)
+        config = cls(**payload)
+        for key in tuple(kwargs):
+            if hasattr(config, key) or key == "name_or_path":
+                setattr(config, key, kwargs.pop(key))
+        return (config, kwargs) if return_unused_kwargs else config
+
+
+__all__ = [
+    "ModilifyMk2Config", "ModilifyMk2TextConfig", "CONFIG_PROTOCOL_ERROR",
+    "COMMIT_READINESS_OBJECTIVE", "COMMIT_SEQUENCE_LAYERS", "DENOISE_TEMPERATURE",
+    "HISTORY_VIEWS", "MEMORY_SCHEME", "MEMORY_ARCHITECTURE", "STATE_SCHEMA_VERSION",
+    "TRAINING_SCHEME", "VOCAB_CHUNK_SIZE",
+    "TOKEN_CE_SUPERVISION", "TOKEN_CE_NORMALIZATION",
+    "require_current_checkpoint_protocol",
+]

continuous_batching.py

@@ -0,0 +1,1737 @@
+"""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
+denoises a transient bidirectional canvas and may accept a ragged token chunk.
+This module consequently owns the request runner while preserving the public
+manager lifecycle and ``GenerationOutput`` contract.
+
+Accepted prefix K/V is stored per request without padding.  Every heavy denoise
+step creates a temporary, left-padded batched cache view.  The decoder only
+reads that view, so padding can never become persistent or evict real tokens
+from a sliding-window cache.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import hashlib
+import math
+import os
+import queue
+import threading
+import time
+import uuid
+import warnings
+from collections import defaultdict, deque
+from collections.abc import Callable, Generator, Sequence
+from dataclasses import asdict, dataclass, field, is_dataclass, replace
+from typing import Any
+
+import torch
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation.configuration_utils import ContinuousBatchingConfig
+from transformers.generation.continuous_batching.requests import (
+    GenerationOutput,
+    RequestStatus,
+)
+
+from .commit_policy import fused_commit_failure_rate, select_commit_lengths
+from .generation_modilify_mk2 import (
+    ModilifyMk2GenerationConfig,
+    ModilifyMk2GenerationOutput,
+    ModilifyMk2RollingState,
+    NoiseCanvasSampler,
+    _add_repetition_history,
+    _flatten_token_ids,
+    build_denoise_trace_event,
+)
+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(
+    {
+        "turn_end",
+        "eos",
+        "max_new_tokens",
+        "max_denoising_steps",
+        "episode_watchdog",
+        "cancelled",
+        "error",
+    }
+)
+
+
+def continuous_config_fingerprint(
+    generation_config: Any,
+    continuous_batching_config: ContinuousBatchingConfig | None,
+) -> str:
+    """Return a stable-enough in-process fingerprint for persistent reuse."""
+
+    generation_payload = (
+        generation_config.to_dict()
+        if hasattr(generation_config, "to_dict")
+        else vars(generation_config)
+    )
+    batching = continuous_batching_config or ContinuousBatchingConfig()
+    batching_payload = asdict(batching) if is_dataclass(batching) else vars(batching)
+    return repr(
+        (
+            sorted(generation_payload.items(), key=lambda item: item[0]),
+            sorted(batching_payload.items(), key=lambda item: item[0]),
+        )
+    )
+
+
+@dataclass
+class ModilifyMk2ContinuousGenerationOutput(GenerationOutput):
+    """Official ``GenerationOutput`` plus ModilifyMk2 request-local diagnostics."""
+
+    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
+    scheduler_run_id: str | None = None
+    queue_seconds: float = 0.0
+    inference_seconds: float = 0.0
+    total_seconds: float = 0.0
+    last_step_batch_size: int = 0
+    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."""
+
+        return self.status in {RequestStatus.FINISHED, RequestStatus.FAILED}
+
+
+@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
+    eos_token_ids: tuple[int, ...]
+    streaming: bool
+    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
+    finished_time: float = -1.0
+    generated_tokens: list[int] = field(default_factory=list)
+    logprobs: list[float] = field(default_factory=list)
+    timestamps: list[float] = field(default_factory=list)
+    cache: Cache | None = None
+    rolling_state: ModilifyMk2RollingState | None = None
+    repetition_history: torch.BoolTensor | None = None
+    generator: torch.Generator | None = None
+    logical_length: int = 0
+    max_iterations: int = 0
+    reserved_blocks: int = 0
+    denoise_steps: int = 0
+    jumps: int = 0
+    forced_jump_tokens: int = 0
+    shifts: int = 0
+    stop_reason: str | None = None
+    error: str | None = None
+    terminal_emitted: bool = False
+    last_step_batch_size: int = 0
+    last_delta_tokens: list[int] = field(default_factory=list)
+
+
+def _clone_tensor_row(value: torch.Tensor, row: int) -> torch.Tensor:
+    return value[row : row + 1].clone()
+
+
+def _slice_rolling_state(state: ModilifyMk2RollingState, row: int) -> ModilifyMk2RollingState:
+    selected = slice(row, row + 1)
+    return ModilifyMk2RollingState(
+        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),
+
+
+        head=_clone_tensor_row(state.head, row),
+    )
+
+
+def _pack_rolling_states(states: Sequence[ModilifyMk2RollingState]) -> ModilifyMk2RollingState:
+    return ModilifyMk2RollingState(
+        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]),
+
+
+        head=torch.cat([state.head for state in states], dim=0),
+    )
+
+
+class ModilifyMk2LogicalCachePool:
+    """Per-request hole-free cache storage with ephemeral batched read views."""
+
+    def __init__(self, model: Any, *, max_batch_tokens: int | None = None) -> None:
+        self.model = model
+        self.text_config = model.config.get_text_config(decoder=True)
+        self.device = model.model.decoder.embed_tokens.weight.device
+        self.max_batch_tokens = max_batch_tokens
+
+    def new_cache(self) -> DynamicCache:
+        return DynamicCache(config=self.text_config)
+
+    @torch.inference_mode()
+    def prefill(self, prompt_ids: Sequence[int]) -> Cache:
+        cache = self.new_cache()
+        chunk_size = self.max_batch_tokens or len(prompt_ids)
+        for start in range(0, len(prompt_ids), chunk_size):
+            stop = min(start + chunk_size, len(prompt_ids))
+            tokens = torch.tensor(
+                [list(prompt_ids[start:stop])], device=self.device, dtype=torch.long
+            )
+            mask = torch.ones(1, stop, device=self.device, dtype=torch.bool)
+            positions = torch.arange(
+                start, stop, device=self.device, dtype=torch.int32
+            ).unsqueeze(0)
+            cache = self.model.model.encoder(
+                input_ids=tokens,
+                attention_mask=mask,
+                past_key_values=cache,
+                position_ids=positions,
+            ).past_key_values
+        return cache
+
+    @torch.inference_mode()
+    def append(self, state: ModilifyMk2RequestState, token_ids: Sequence[int]) -> None:
+        if not token_ids:
+            return
+        if state.cache is None:
+            raise RuntimeError("Cannot append tokens before request prefill.")
+        tokens = torch.tensor([list(token_ids)], device=self.device, dtype=torch.long)
+        positions = torch.arange(
+            state.logical_length,
+            state.logical_length + tokens.shape[1],
+            device=self.device,
+            dtype=torch.int32,
+        ).unsqueeze(0)
+        mask = torch.ones(
+            1,
+            state.logical_length + tokens.shape[1],
+            device=self.device,
+            dtype=torch.bool,
+        )
+        state.cache = self.model.model.encoder(
+            input_ids=tokens,
+            attention_mask=mask,
+            past_key_values=state.cache,
+            position_ids=positions,
+        ).past_key_values
+
+    def pack(
+        self, states: Sequence[ModilifyMk2RequestState]
+    ) -> tuple[DynamicCache, torch.BoolTensor, torch.LongTensor]:
+        if not states or any(state.cache is None for state in states):
+            raise ValueError("Every packed request must have an initialized cache.")
+        logical_lengths = torch.tensor(
+            [state.logical_length for state in states],
+            device=self.device,
+            dtype=torch.long,
+        )
+        maximum_length = int(logical_lengths.max())
+        attention_mask = torch.arange(
+            maximum_length, device=self.device
+        )[None, :].ge(maximum_length - logical_lengths[:, None])
+
+        packed = self.new_cache()
+        source_caches = [state.cache for state in states]
+        assert all(cache is not None for cache in source_caches)
+        if any(len(cache.layers) != len(packed.layers) for cache in source_caches):
+            raise RuntimeError("Request cache layer structures differ.")
+
+        for layer_index, packed_layer in enumerate(packed.layers):
+            source_layers = [cache.layers[layer_index] for cache in source_caches]
+            if any(not layer.is_initialized for layer in source_layers):
+                raise RuntimeError("Request cache contains an uninitialized layer.")
+            stored_lengths = [int(layer.keys.shape[-2]) for layer in source_layers]
+            maximum_stored = max(stored_lengths)
+
+            def padded(name: str) -> torch.Tensor:
+                values = []
+                for layer, stored_length in zip(source_layers, stored_lengths, strict=True):
+                    value = getattr(layer, name)
+                    if stored_length < maximum_stored:
+                        padding = value.new_zeros(
+                            value.shape[0],
+                            value.shape[1],
+                            maximum_stored - stored_length,
+                            value.shape[3],
+                        )
+                        value = torch.cat((padding, value), dim=-2)
+                    values.append(value)
+                return torch.cat(values, dim=0)
+
+            keys = padded("keys")
+            values = padded("values")
+            packed_layer.lazy_initialization(keys, values)
+            packed_layer.keys = keys
+            packed_layer.values = values
+            if hasattr(packed_layer, "cumulative_length"):
+                packed_layer.cumulative_length = maximum_length
+        return packed, attention_mask, logical_lengths
+
+
+class ModilifyMk2ContinuousBatchingManager:
+    """FIFO/prefill-first continuous manager compatible with Transformers APIs."""
+
+    def __init__(
+        self,
+        model: Any,
+        generation_config: ModilifyMk2GenerationConfig | None,
+        continuous_batching_config: ContinuousBatchingConfig | None,
+        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)
+            or ModilifyMk2GenerationConfig.from_model_config(model.config)
+        )
+        if not isinstance(self.generation_config, ModilifyMk2GenerationConfig):
+            payload = self.generation_config.to_dict()
+            self.generation_config = ModilifyMk2GenerationConfig(**payload)
+        self.continuous_batching_config = copy.deepcopy(
+            continuous_batching_config or ContinuousBatchingConfig()
+        )
+        self.config_fingerprint = continuous_config_fingerprint(
+            self.generation_config, self.continuous_batching_config
+        )
+        self._validate_config()
+
+        self.device = model.model.decoder.embed_tokens.weight.device
+        self.dtype = model.model.decoder.embed_tokens.weight.dtype
+        self.cache_pool = ModilifyMk2LogicalCachePool(
+            model,
+            max_batch_tokens=self.continuous_batching_config.max_batch_tokens,
+        )
+        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
+
+        configured_requests = self.continuous_batching_config.max_requests_per_batch
+        self.max_requests_per_batch = int(configured_requests or 8)
+        max_batch_tokens = self.continuous_batching_config.max_batch_tokens
+        if max_batch_tokens is not None:
+            token_capacity = int(max_batch_tokens) // int(model.config.canvas_length)
+            if token_capacity < 1:
+                raise ValueError(
+                    "`max_batch_tokens` must fit at least one ModilifyMk2 canvas."
+                )
+            self.max_requests_per_batch = min(
+                self.max_requests_per_batch, token_capacity
+            )
+        self.block_size = int(self.continuous_batching_config.block_size)
+        self.block_capacity = self._resolve_block_capacity()
+        self._base_seed = (
+            int(self.continuous_batching_config.seed)
+            if self.continuous_batching_config.seed is not None
+            else int(torch.initial_seed())
+        )
+
+        self._condition = threading.Condition(threading.RLock())
+        self._pending: deque[ModilifyMk2RequestState] = deque()
+        self._active: dict[str, ModilifyMk2RequestState] = {}
+        self._known_request_ids: set[str] = set()
+        self._cancelled: set[str] = set()
+        self._output_queue: queue.Queue[ModilifyMk2ContinuousGenerationOutput] = queue.Queue()
+        self._stashed_outputs: dict[
+            str, deque[ModilifyMk2ContinuousGenerationOutput]
+        ] = defaultdict(deque)
+        self._result_handlers: dict[str, tuple[Callable, asyncio.AbstractEventLoop]] = {}
+        self._thread: threading.Thread | None = None
+        self._finished = threading.Event()
+        self.fatal_error: BaseException | None = None
+        self._input_closed = False
+        self._hard_stop = False
+        self._keep_for_next_session = False
+        self._request_counter = 0
+        self._active_reserved_blocks = 0
+        self._stats = {
+            "submitted": 0,
+            "admitted": 0,
+            "completed": 0,
+            "failed": 0,
+            "cancelled": 0,
+            "model_steps": 0,
+            "generated_tokens": 0,
+            "max_observed_batch_size": 0,
+            "peak_reserved_blocks": 0,
+            "peak_cache_blocks": 0,
+            "active_slot_steps": 0,
+            "slot_capacity_steps": 0,
+        }
+
+        turn_end = self.generation_config.turn_end_token_id
+        self.turn_end_token_id = int(
+            model.config.turn_end_token_id if turn_end is None else turn_end
+        )
+        self.repetition_penalty = float(self.generation_config.repetition_penalty)
+        self.excluded_repetition_token_ids = _flatten_token_ids(
+            self.generation_config.repetition_penalty_exclude_token_ids,
+            self.generation_config.pad_token_id,
+            self.generation_config.bos_token_id,
+            self.generation_config.eos_token_id,
+            self.generation_config.turn_end_token_id,
+            getattr(model.config, "image_token_id", None),
+        )
+
+    def _validate_config(self) -> None:
+        config = self.continuous_batching_config
+        positive_optional = (
+            "num_blocks",
+            "max_batch_tokens",
+            "max_requests_per_batch",
+        )
+        if not isinstance(config.block_size, int) or config.block_size < 4:
+            raise ValueError("`block_size` must be an integer greater than or equal to 4.")
+        for name in positive_optional:
+            value = getattr(config, name)
+            if value is not None and (not isinstance(value, int) or value <= 0):
+                raise ValueError(f"`{name}` must be a positive integer when set.")
+        if config.max_blocks_per_request is not None and (
+            not isinstance(config.max_blocks_per_request, int)
+            or config.max_blocks_per_request < 0
+        ):
+            raise ValueError("`max_blocks_per_request` must be a non-negative integer.")
+        if not isinstance(config.max_queue_size, int) or config.max_queue_size < 0:
+            raise ValueError("`max_queue_size` must be a non-negative integer.")
+        if config.scheduler_type not in {"fifo", "prefill_first"}:
+            raise ValueError("ModilifyMk2 continuous batching supports `fifo` and `prefill_first`.")
+        if config.max_memory_percent is not None and not (
+            0.0 < float(config.max_memory_percent) <= 1.0
+        ):
+            raise ValueError("`max_memory_percent` must be in (0, 1].")
+        if config.use_async_batching is True:
+            raise ValueError("ModilifyMk2 continuous batching currently uses synchronous model steps.")
+        requested_graphs = config.use_cuda_graph
+        if requested_graphs is True or (
+            isinstance(requested_graphs, tuple) and any(requested_graphs)
+        ):
+            raise ValueError("CUDA graphs are not supported by the ragged ModilifyMk2 runner.")
+        if config.cpu_offload_space is not None and config.cpu_offload_space > 0:
+            raise ValueError("CPU cache offload is not supported by the ModilifyMk2 runner.")
+        if int(config.default_compile_level or 0) > 0:
+            raise ValueError("Continuous ModilifyMk2 compilation is not supported yet.")
+        if config.varlen_compile_config is not None or config.decode_compile_config is not None:
+            raise ValueError("Continuous ModilifyMk2 compilation is not supported yet.")
+        if config.use_default_compile_configs is True:
+            raise ValueError("Continuous ModilifyMk2 compilation is not supported yet.")
+        if int(config.q_padding_interval_size or 0) > 0 or int(
+            config.kv_padding_interval_size or 0
+        ) > 0:
+            raise ValueError("Compiled continuous padding intervals are not supported.")
+        if config.max_cached_graphs is not None:
+            raise ValueError("Cached continuous graphs are not supported.")
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            if torch.distributed.get_world_size() > 1:
+                raise ValueError(
+                    "Tensor/distributed parallel continuous batching is not supported."
+                )
+        if getattr(self.model, "device_mesh", None) is not None or getattr(
+            self.model, "_device_mesh", None
+        ) is not None:
+            raise ValueError("Tensor-parallel continuous batching is not supported.")
+        # Prefix sharing would make request ownership and row-local RNG/state
+        # ambiguous.  Normalize this optimization off rather than silently use it.
+        config.allow_block_sharing = False
+
+    def _available_memory_bytes(self) -> int | None:
+        if self.device.type == "cuda" and torch.cuda.is_available():
+            free, _ = torch.cuda.mem_get_info(self.device)
+            return int(free)
+        if self.device.type == "mps" and torch.backends.mps.is_available():
+            return max(
+                0,
+                int(torch.mps.recommended_max_memory())
+                - int(torch.mps.driver_allocated_memory()),
+            )
+        if self.device.type == "cpu":
+            try:
+                import psutil
+
+                return int(psutil.virtual_memory().available)
+            except (ImportError, OSError, ValueError):
+                pass
+            try:
+                return int(os.sysconf("SC_AVPHYS_PAGES")) * int(
+                    os.sysconf("SC_PAGE_SIZE")
+                )
+            except (OSError, TypeError, ValueError):
+                return None
+        return None
+
+    def _estimated_block_bytes(self) -> int:
+        config = self.model.config.text_config
+        layer_types = list(config.layer_types)
+        local_heads = int(config.num_key_value_heads)
+        local_dim = int(config.head_dim)
+        global_heads = int(
+            getattr(config, "num_global_key_value_heads", None) or local_heads
+        )
+        global_dim = int(getattr(config, "global_head_dim", None) or local_dim)
+        per_token = 0
+        for layer_type in layer_types:
+            if layer_type == "full_attention":
+                heads, dimension = global_heads, global_dim
+            else:
+                heads, dimension = local_heads, local_dim
+            per_token += 2 * heads * dimension * torch.empty((), dtype=self.dtype).element_size()
+        return max(1, per_token * int(self.continuous_batching_config.block_size))
+
+    def _resolve_block_capacity(self) -> int | None:
+        capacity = self.continuous_batching_config.num_blocks
+        percent = self.continuous_batching_config.max_memory_percent
+        available = self._available_memory_bytes()
+        if percent is None and capacity is None:
+            # Never make the default cache silently unbounded. This fraction is
+            # applied to currently available device/host memory after model load.
+            percent = 0.8
+        if percent is not None and available is None:
+            raise RuntimeError(
+                "Cannot infer available cache memory on this device; set `num_blocks` "
+                "explicitly instead of `max_memory_percent`."
+            )
+        if percent is not None and available is not None:
+            memory_blocks = int(
+                available * float(percent) / self._estimated_block_bytes()
+            )
+            capacity = memory_blocks if capacity is None else min(int(capacity), memory_blocks)
+        return None if capacity is None else max(0, int(capacity))
+
+    @staticmethod
+    def _block_footprint(reservations: Sequence[int]) -> int:
+        """Return persistent plus temporary packed-cache block equivalents."""
+
+        if not reservations:
+            return 0
+        return sum(reservations) + len(reservations) * max(reservations)
+
+    def _current_block_footprint(self) -> int:
+        return self._block_footprint(
+            [state.reserved_blocks for state in self._active.values()]
+        )
+
+    def _derive_seed(self, request_id: str) -> int:
+        digest = hashlib.sha256(
+            str(self._base_seed).encode("ascii")
+            + b"\0"
+            + request_id.encode("utf-8")
+        ).digest()
+        return int.from_bytes(digest[:8], "big") & ((1 << 63) - 1)
+
+    @property
+    def stats(self) -> dict[str, Any]:
+        with self._condition:
+            capacity_steps = int(self._stats["slot_capacity_steps"])
+            return {
+                "scheduler_run_id": self.run_id,
+                **self._stats,
+                "slot_utilization": (
+                    float(self._stats["active_slot_steps"]) / capacity_steps
+                    if capacity_steps
+                    else 0.0
+                ),
+                "active_requests": len(self._active),
+                "pending_requests": len(self._pending),
+                "max_requests_per_batch": self.max_requests_per_batch,
+                "block_capacity": -1 if self.block_capacity is None else self.block_capacity,
+                "reserved_blocks": self._active_reserved_blocks,
+                "cache_blocks": self._current_block_footprint(),
+            }
+
+    def is_running(self) -> bool:
+        return self._thread is not None and self._thread.is_alive()
+
+    def warmup(self) -> None:
+        if self.destroyed:
+            raise RuntimeError("Cannot warm up a destroyed manager.")
+        # CUDA graphs and static-shape compilation are intentionally unsupported;
+        # normal eager kernels warm naturally on the first real batch.
+        self.warmed_up = True
+
+    def start(self) -> None:
+        if self._keep_for_next_session:
+            self._prepare_for_next_session()
+        with self._condition:
+            if self.destroyed:
+                raise RuntimeError("Cannot start a destroyed manager.")
+            if self.is_running():
+                return
+            self._finished.clear()
+            self.fatal_error = None
+            self._hard_stop = False
+            self._thread = threading.Thread(
+                target=self._run_generation_loop,
+                name=f"modilify_mk2-continuous-{self.run_id[:8]}",
+                daemon=True,
+            )
+            self._thread.start()
+
+    def join(
+        self,
+        stop_trigger_time: float | None = None,
+        timeout: float | None = None,
+    ) -> None:
+        """Wait for the current worker, matching the official manager lifecycle."""
+
+        del stop_trigger_time
+        with self._condition:
+            thread = self._thread
+        if thread is None or thread is threading.current_thread():
+            return
+        thread.join(timeout=timeout)
+        if thread.is_alive():
+            raise TimeoutError("Timed out waiting for continuous generation to stop.")
+
+    def _prepare_for_next_session(self) -> None:
+        """Finish an asynchronous prior stop and reopen a cached manager safely."""
+
+        with self._condition:
+            if not self._keep_for_next_session:
+                return
+            thread = self._thread
+        if thread is not None and thread.is_alive():
+            thread.join()
+        with self._condition:
+            if self.destroyed:
+                raise RuntimeError("Cannot reuse a destroyed manager.")
+            if self._pending or self._active:
+                raise RuntimeError("Cannot reuse a manager with unfinished requests.")
+            self._input_closed = False
+            self._hard_stop = False
+            self._keep_for_next_session = False
+            self.fatal_error = None
+            self._cancelled.clear()
+            self._condition.notify_all()
+
+    def close_input(self) -> None:
+        """Stop accepting requests and let the iterator drain all submitted work."""
+
+        with self._condition:
+            self._input_closed = True
+            self._condition.notify_all()
+
+    def stop(
+        self,
+        block: bool = True,
+        timeout: float | None = None,
+        keep_for_next_session: bool = False,
+        hard_stop: bool = False,
+    ) -> None:
+        with self._condition:
+            self._input_closed = True
+            self._hard_stop = bool(hard_stop)
+            self._keep_for_next_session = bool(keep_for_next_session)
+            if hard_stop:
+                self._cancelled.update(self._known_request_ids)
+            self._condition.notify_all()
+            thread = self._thread
+        if hard_stop and (thread is None or not thread.is_alive()):
+            self._apply_cancellations()
+        if block and thread is not None:
+            self.join(timeout=timeout)
+        if keep_for_next_session and not self.is_running():
+            with self._condition:
+                self._input_closed = False
+                self._hard_stop = False
+                self._keep_for_next_session = False
+                self.fatal_error = None
+
+    def destroy(self) -> None:
+        if self.destroyed:
+            return
+        self.stop(block=True, hard_stop=True)
+        self.destroyed = True
+        with self._condition:
+            self._pending.clear()
+            self._active.clear()
+            self._condition.notify_all()
+
+    def add_request(
+        self,
+        input_ids: list[int],
+        request_id: str | None = None,
+        max_new_tokens: int | None = None,
+        streaming: bool = False,
+        record_timestamps: bool = False,
+        eos_token_id: int | list[int] | None = None,
+        **request_kwargs: Any,
+    ) -> str:
+        if not input_ids or any(
+            not isinstance(token_id, int) or isinstance(token_id, bool)
+            for token_id in input_ids
+        ):
+            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.")
+        if max_denoising_steps is not None and (
+            not isinstance(max_denoising_steps, int) or max_denoising_steps <= 0
+        ):
+            raise ValueError("`max_denoising_steps` must be a positive integer when set.")
+
+        with self._condition:
+            if self.destroyed or self._input_closed:
+                raise RuntimeError("Continuous batching manager is not accepting requests.")
+            if self.fatal_error is not None:
+                raise RuntimeError("Continuous batching manager has failed.") from self.fatal_error
+            if request_id is None:
+                request_id = f"req_{self._request_counter}"
+                self._request_counter += 1
+            if request_id in self._known_request_ids:
+                raise ValueError(f"Duplicate continuous request ID: {request_id}")
+            queue_limit = int(self.continuous_batching_config.max_queue_size)
+            deadline = time.monotonic() + 10.0
+            while queue_limit and len(self._pending) >= queue_limit:
+                if not self.is_running():
+                    raise queue.Full(
+                        "Continuous request queue is full; start the manager before "
+                        "submitting more requests."
+                    )
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise queue.Full("Continuous request queue remained full for 10 seconds.")
+                self._condition.wait(timeout=remaining)
+                if self.destroyed or self._input_closed:
+                    raise RuntimeError(
+                        "Continuous batching manager stopped while waiting for queue space."
+                    )
+                if self.fatal_error is not None:
+                    raise RuntimeError("Continuous batching manager has failed.") from self.fatal_error
+            # The worker can close/fail the manager while this producer is
+            # asleep. Recheck under the same lock immediately before append.
+            if self.destroyed or self._input_closed:
+                raise RuntimeError("Continuous batching manager is not accepting requests.")
+            if self.fatal_error is not None:
+                raise RuntimeError("Continuous batching manager has failed.") from self.fatal_error
+            if request_id in self._known_request_ids:
+                raise ValueError(f"Duplicate continuous request ID: {request_id}")
+            configured_eos = self.generation_config.eos_token_id if eos_token_id is None else eos_token_id
+            if configured_eos is None:
+                configured_eos = self.model.config.eos_token_id
+            eos_values = (
+                [configured_eos]
+                if isinstance(configured_eos, int)
+                else list(configured_eos or [])
+            )
+            stop_ids = tuple(
+                dict.fromkeys(
+                    [self.turn_end_token_id, *(int(value) for value in eos_values if int(value) >= 0)]
+                )
+            )
+            resolved_seed = self._derive_seed(request_id) if seed is None else int(seed)
+            state = ModilifyMk2RequestState(
+                request_id=request_id,
+                prompt_ids=list(input_ids),
+                max_new_tokens=int(limit),
+                eos_token_ids=stop_ids,
+                streaming=bool(streaming),
+                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
+            )
+            self._pending.append(state)
+            self._known_request_ids.add(request_id)
+            self._stats["submitted"] += 1
+            self._condition.notify_all()
+            return request_id
+
+    def add_requests(
+        self,
+        inputs: list[list[int]],
+        max_new_tokens: int | None = None,
+        streaming: bool = False,
+        record_timestamps: bool = False,
+        **request_kwargs: Any,
+    ) -> list[str]:
+        request_ids = request_kwargs.pop("request_ids", None)
+        seeds = request_kwargs.pop("seeds", None)
+        if request_ids is not None and len(request_ids) != len(inputs):
+            raise ValueError("`request_ids` must contain one ID per request.")
+        if seeds is not None and len(seeds) != len(inputs):
+            raise ValueError("`seeds` must contain one seed per request.")
+        result = []
+        for index, input_ids in enumerate(inputs):
+            per_request = dict(request_kwargs)
+            if seeds is not None:
+                per_request["seed"] = seeds[index]
+            result.append(
+                self.add_request(
+                    input_ids=input_ids,
+                    request_id=None if request_ids is None else request_ids[index],
+                    max_new_tokens=max_new_tokens,
+                    streaming=streaming,
+                    record_timestamps=record_timestamps,
+                    **per_request,
+                )
+            )
+        return result
+
+    def cancel_request(self, request_id: str) -> None:
+        with self._condition:
+            if request_id in self._known_request_ids:
+                self._cancelled.add(request_id)
+                self._condition.notify_all()
+
+    def register_result_handler(self, request_id: str, callback: Callable) -> None:
+        loop = asyncio.get_running_loop()
+        with self._condition:
+            self._result_handlers[request_id] = (callback, loop)
+
+    def _pop_stashed(self, request_id: str | None):
+        with self._condition:
+            if request_id is not None:
+                values = self._stashed_outputs.get(request_id)
+                if values:
+                    return values.popleft()
+                return None
+            for values in self._stashed_outputs.values():
+                if values:
+                    return values.popleft()
+        return None
+
+    def _has_stashed_outputs(self) -> bool:
+        with self._condition:
+            return any(values for values in self._stashed_outputs.values())
+
+    def get_result(
+        self, request_id: str | None = None, timeout: float | None = None
+    ) -> ModilifyMk2ContinuousGenerationOutput | None:
+        stashed = self._pop_stashed(request_id)
+        if stashed is not None:
+            return stashed
+        if not self.is_running() and self._output_queue.empty():
+            return None
+        deadline = None if timeout is None else time.monotonic() + timeout
+        while True:
+            remaining = None if deadline is None else max(0.0, deadline - time.monotonic())
+            if remaining == 0.0:
+                return None
+            try:
+                output = self._output_queue.get(timeout=remaining)
+            except queue.Empty:
+                return None
+            if request_id is None or output.request_id == request_id:
+                return output
+            with self._condition:
+                self._stashed_outputs[output.request_id].append(output)
+
+    def __iter__(self) -> Generator[ModilifyMk2ContinuousGenerationOutput, None, None]:
+        while True:
+            output = self.get_result(timeout=0.05)
+            if output is not None:
+                yield output
+                continue
+            if self._finished.is_set() and self._output_queue.empty():
+                if not self._has_stashed_outputs():
+                    return
+
+    def request_id_iter(
+        self, request_id: str
+    ) -> Generator[ModilifyMk2ContinuousGenerationOutput, None, None]:
+        while True:
+            output = self.get_result(request_id=request_id, timeout=0.05)
+            if output is not None:
+                yield output
+                if output.is_finished():
+                    return
+            elif self._finished.is_set():
+                return
+
+    def _deliver(self, output: ModilifyMk2ContinuousGenerationOutput) -> None:
+        handler = None
+        with self._condition:
+            handler = self._result_handlers.get(output.request_id)
+            if output.is_finished():
+                self._result_handlers.pop(output.request_id, None)
+        if handler is None:
+            self._output_queue.put(output)
+        else:
+            callback, loop = handler
+            try:
+                loop.call_soon_threadsafe(callback, output)
+            except RuntimeError as error:
+                # A callback owner may close its event loop while a terminal
+                # event is in flight. Preserve the result for pull consumers
+                # instead of turning that client race into a worker fatality.
+                warnings.warn(
+                    f"Result callback loop closed for {output.request_id}: {error!r}",
+                    stacklevel=2,
+                )
+                self._output_queue.put(output)
+
+    def _output_for(
+        self,
+        state: ModilifyMk2RequestState,
+        *,
+        stream_update: bool = False,
+        delta_tokens: Sequence[int] | None = None,
+    ) -> ModilifyMk2ContinuousGenerationOutput:
+        now = time.perf_counter()
+        finished = state.status in {RequestStatus.FINISHED, RequestStatus.FAILED}
+        end = state.finished_time if finished else -1.0
+        shifts = max(1, state.shifts)
+        steps = max(1, state.denoise_steps)
+        return ModilifyMk2ContinuousGenerationOutput(
+            request_id=state.request_id,
+            prompt_ids=list(state.prompt_ids),
+            generated_tokens=list(state.generated_tokens),
+            logprobs=list(state.logprobs),
+            error=state.error,
+            status=state.status,
+            created_time=state.created_time,
+            lifespan=(state.started_time, end),
+            timestamps=(list(state.timestamps) if state.record_timestamps else None),
+            stop_reason=state.stop_reason,
+            committed_tokens=len(state.generated_tokens),
+            denoise_steps=state.denoise_steps,
+            no_progress_steps=(
+                0
+                if state.rolling_state is None
+                else int(state.rolling_state.latent_state.stagnation_steps[0])
+            ),
+            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,
+            scheduler_run_id=self.run_id,
+            queue_seconds=max(0.0, state.started_time - state.created_time),
+            inference_seconds=(
+                max(0.0, (end if finished else now) - state.started_time)
+                if state.started_time >= 0
+                else 0.0
+            ),
+            total_seconds=max(0.0, (end if finished else now) - state.created_time),
+            last_step_batch_size=state.last_step_batch_size,
+            is_stream_update=stream_update,
+            delta_tokens=list(
+                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(
+        self,
+        state: ModilifyMk2RequestState,
+        reason: str,
+        error: BaseException | str | None = None,
+    ) -> None:
+        if state.terminal_emitted:
+            return
+        if reason not in _TERMINAL_REASONS:
+            raise ValueError(f"Unknown continuous stop reason: {reason}")
+        state.stop_reason = reason
+        state.error = None if error is None else (str(error) if isinstance(error, str) else repr(error))
+        state.status = (
+            RequestStatus.FAILED
+            if reason in {"cancelled", "error"} or error is not None
+            else RequestStatus.FINISHED
+        )
+        state.finished_time = time.perf_counter()
+        state.terminal_emitted = True
+        self._stats["completed"] += 1
+        if reason == "cancelled":
+            self._stats["cancelled"] += 1
+        elif error is not None:
+            self._stats["failed"] += 1
+        self._deliver(self._output_for(state))
+
+    def _fail_all_requests(self, error: BaseException) -> None:
+        """Convert an unexpected worker failure into one terminal result per request."""
+
+        self.fatal_error = error
+        with self._condition:
+            pending = list(self._pending)
+            active = list(self._active.values())
+            self._pending.clear()
+            self._active.clear()
+            self._active_reserved_blocks = 0
+            self._input_closed = True
+            for state in [*active, *pending]:
+                self._finish(state, "error", error)
+            self._condition.notify_all()
+
+    def _request_fits(self, state: ModilifyMk2RequestState) -> bool:
+        per_request_limit = self.continuous_batching_config.max_blocks_per_request
+        if per_request_limit not in (None, 0) and state.reserved_blocks > per_request_limit:
+            return False
+        if self.block_capacity is None:
+            return True
+        reservations = [
+            *(active.reserved_blocks for active in self._active.values()),
+            state.reserved_blocks,
+        ]
+        return self._block_footprint(reservations) <= self.block_capacity
+
+    def _request_can_ever_fit(self, state: ModilifyMk2RequestState) -> bool:
+        per_request_limit = self.continuous_batching_config.max_blocks_per_request
+        if per_request_limit not in (None, 0) and state.reserved_blocks > per_request_limit:
+            return False
+        return (
+            self.block_capacity is None
+            or self._block_footprint([state.reserved_blocks]) <= self.block_capacity
+        )
+
+    def _initialize_request(self, state: ModilifyMk2RequestState) -> None:
+        generator = torch.Generator(device=self.device)
+        generator.manual_seed(state.seed)
+        state.generator = generator
+        canvas = self.sampler.initialize_canvas(
+            1, self.device, generators=[generator]
+        )
+        canvas[:, state.max_new_tokens:] = self.pad_token_id
+        dtype = self.model.model.decoder.embed_tokens.weight.dtype
+        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,
+            enable_gdn2=True,
+        )
+        state.rolling_state = ModilifyMk2RollingState(
+            canvas=canvas,
+            confidence=torch.zeros(1, canvas_length, device=self.device, dtype=torch.float32),
+            entropy=torch.full(
+                (1, canvas_length),
+                math.log(self.model.config.text_config.vocab_size),
+                device=self.device,
+                dtype=torch.float32,
+            ),
+            age=torch.zeros(1, canvas_length, device=self.device, dtype=torch.int32),
+            latent_state=latent,
+
+
+            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)
+        if self.repetition_penalty != 1.0:
+            state.repetition_history = torch.zeros(
+                self.model.config.text_config.vocab_size,
+                device=self.device,
+                dtype=torch.bool,
+            )
+            prompt = torch.tensor([state.prompt_ids], device=self.device, dtype=torch.long)
+            _add_repetition_history(
+                state.repetition_history.unsqueeze(0),
+                prompt,
+                torch.ones_like(prompt, dtype=torch.bool),
+                self.excluded_repetition_token_ids,
+            )
+        state.max_iterations = deterministic_episode_iteration_bound(
+            torch.tensor([state.max_new_tokens]),
+            max_ponder_steps=self.generation_config.max_ponder_steps,
+        )
+        state.started_time = time.perf_counter()
+        state.status = RequestStatus.DECODING
+
+    def _apply_cancellations(self) -> None:
+        with self._condition:
+            cancelled = set(self._cancelled)
+            self._cancelled.clear()
+            if not cancelled:
+                return
+            retained = deque()
+            while self._pending:
+                state = self._pending.popleft()
+                if state.request_id in cancelled:
+                    self._finish(state, "cancelled", "request cancelled")
+                else:
+                    retained.append(state)
+            self._pending = retained
+            for request_id in cancelled:
+                state = self._active.pop(request_id, None)
+                if state is not None:
+                    self._active_reserved_blocks -= state.reserved_blocks
+                    self._finish(state, "cancelled", "request cancelled")
+            self._condition.notify_all()
+
+    def _admit_requests(self) -> None:
+        while True:
+            with self._condition:
+                if len(self._active) >= self.max_requests_per_batch or not self._pending:
+                    return
+                state = self._pending[0]
+                if not self._request_can_ever_fit(state):
+                    self._pending.popleft()
+                    self._finish(
+                        state,
+                        "error",
+                        "request exceeds continuous cache block limits",
+                    )
+                    self._condition.notify_all()
+                    continue
+                if not self._request_fits(state):
+                    return
+                self._pending.popleft()
+                self._condition.notify_all()
+            try:
+                self._initialize_request(state)
+            except Exception as error:
+                self._finish(state, "error", error)
+                continue
+            with self._condition:
+                if state.request_id in self._cancelled:
+                    self._cancelled.remove(state.request_id)
+                    self._finish(state, "cancelled", "request cancelled")
+                    continue
+                self._active[state.request_id] = state
+                self._active_reserved_blocks += state.reserved_blocks
+                self._stats["admitted"] += 1
+                self._stats["peak_reserved_blocks"] = max(
+                    self._stats["peak_reserved_blocks"],
+                    self._active_reserved_blocks,
+                )
+                self._stats["peak_cache_blocks"] = max(
+                    self._stats["peak_cache_blocks"],
+                    self._current_block_footprint(),
+                )
+                self._condition.notify_all()
+
+    def _select_rowwise_policy(
+        self,
+        states: Sequence[ModilifyMk2RequestState],
+        proposal: torch.LongTensor,
+        normal_failure_rate: torch.Tensor,
+        previous_failure_rate: torch.Tensor,
+        greedy_proposal: torch.LongTensor,
+        jump_failure_rate: torch.Tensor,
+        rolling: ModilifyMk2RollingState,
+    ):
+        decisions = []
+        for row, state in enumerate(states):
+            remaining = state.max_new_tokens - len(state.generated_tokens)
+            decisions.append(
+                select_commit_lengths(
+                    sampled_token_ids=proposal[row : row + 1],
+                    normal_failure_rate=normal_failure_rate[row : row + 1],
+                    previous_failure_rate=previous_failure_rate[row : row + 1],
+                    greedy_token_ids=greedy_proposal[row : row + 1],
+                    jump_failure_rate=jump_failure_rate[row : row + 1],
+                    ponder_steps=rolling.latent_state.ponder_steps[row : row + 1],
+                    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.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,
+                    min_progress=self.generation_config.min_trajectory_progress,
+                )
+            )
+        return (
+            torch.cat([decision.normal_lengths for decision in decisions]),
+            torch.cat([decision.commit_lengths for decision in decisions]),
+            torch.cat([decision.commit_token_ids for decision in decisions]),
+            torch.cat([decision.jump_rows for decision in decisions]),
+            torch.cat([decision.ponder_steps for decision in decisions]),
+            torch.cat([decision.stagnation_steps for decision in decisions]),
+        )
+
+    @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):
+            raise RuntimeError("Active request has no rolling state.")
+        rolling = _pack_rolling_states(rolling_states)  # type: ignore[arg-type]
+        packed_cache, cache_mask, logical_lengths = self.cache_pool.pack(states)
+        batch_size = len(states)
+        canvas_length = int(self.model.config.canvas_length)
+        decoder_positions = (
+            logical_lengths[:, None]
+            + torch.arange(canvas_length, device=self.device)[None, :]
+        ).to(torch.int32)
+        decoder_mask = torch.cat(
+            (
+                cache_mask,
+                torch.ones(
+                    batch_size,
+                    canvas_length,
+                    device=self.device,
+                    dtype=torch.bool,
+                ),
+            ),
+            dim=-1,
+        )
+        repetition_history = None
+        if self.repetition_penalty != 1.0:
+            repetition_history = torch.stack(
+                [state.repetition_history for state in states], dim=0  # type: ignore[list-item]
+            )
+        generators = [state.generator for state in states]
+        if any(generator is None for generator in generators):
+            raise RuntimeError("Active request has no sampling generator.")
+
+        output = self.model(
+            input_ids=None,
+            past_key_values=packed_cache,
+            decoder_input_ids=rolling.canvas,
+            previous_confidence=rolling.confidence,
+            previous_entropy=rolling.entropy,
+            token_age=rolling.age,
+            latent_state=rolling.latent_state,
+
+
+            decoder_position_ids=decoder_positions,
+            decoder_read_cache=True,
+            decoder_attention_mask=decoder_mask,
+            compact_vocab=True,
+            repetition_token_mask=repetition_history,
+            repetition_penalty=self.repetition_penalty,
+            sampling_generators=generators,
+        )
+        required = (
+            output.proposal,
+            output.proposal_confidence,
+            output.token_entropy,
+            output.greedy_proposal,
+            output.greedy_confidence,
+            output.next_latent_state,
+        )
+        if any(value is None for value in required):
+            raise RuntimeError("Compact ModilifyMk2 forward did not return proposal state.")
+        proposal = output.proposal
+        proposal_confidence = output.proposal_confidence
+        token_entropy = output.token_entropy
+        greedy_proposal = output.greedy_proposal
+        greedy_confidence = output.greedy_confidence
+        next_canvas = proposal.clone()
+        next_confidence = proposal_confidence.float()
+        next_latent = replace(
+            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,
+        )
+        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]
+        )
+        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,
+
+
+            head=rolling.head,
+        )
+        normal_failure_rate = fused_commit_failure_rate(
+            proposal_confidence,
+            token_entropy,
+            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),
+            vocab_size=self.model.config.text_config.vocab_size,
+        )
+        jump_failure_rate = fused_commit_failure_rate(
+            greedy_confidence,
+            token_entropy,
+            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),
+            vocab_size=self.model.config.text_config.vocab_size,
+        )
+        previous_failure_rate = fused_commit_failure_rate(
+            rolling.confidence,
+            rolling.entropy,
+            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),
+            vocab_size=self.model.config.text_config.vocab_size,
+        )
+        (
+            normal_commit,
+            commit_lengths,
+            commit_token_ids,
+            jump_rows,
+            next_ponder,
+            next_stagnation,
+        ) = self._select_rowwise_policy(
+            states,
+            proposal,
+            normal_failure_rate,
+            previous_failure_rate,
+            greedy_proposal,
+            jump_failure_rate,
+            rolling,
+        )
+        positions = torch.arange(canvas_length, device=self.device)[None, :]
+        commit_positions = positions.lt(commit_lengths[:, None])
+        policy_prefix_mask = positions.lt(normal_commit[:, None])
+        if bool(jump_rows.any()):
+            next_state = replace(
+                next_state,
+                canvas=torch.where(
+                    commit_positions & jump_rows[:, None],
+                    commit_token_ids,
+                    next_state.canvas,
+                ),
+            )
+        next_state = replace(
+            next_state,
+            latent_state=replace(
+                next_state.latent_state,
+                ponder_steps=next_ponder,
+                stagnation_steps=next_stagnation,
+            ),
+        )
+        unshifted_trace_states = [
+            _slice_rolling_state(next_state, row) for row in range(batch_size)
+        ]
+        if output.working_state is None:
+            raise RuntimeError("Forward did not return working trajectory features.")
+        next_state = self.model._write_committed_memory(
+
+            next_state=next_state,
+            working_state=output.working_state,
+
+            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(
+                commit_lengths,
+                jump_rows=jump_rows,
+                commit_token_ids=commit_token_ids,
+                terminal_token_ids=getattr(
+                    self.model.config, "terminal_token_ids", ()
+                ),
+            ),
+        )
+        shifted = self.model._shift_state_rows(
+            next_state,
+            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)
+        ]
+        selected_confidence = torch.where(
+            jump_rows[:, None], greedy_confidence, proposal_confidence
+        ).float()
+
+        finished_ids = []
+        for row, state in enumerate(states):
+            state.last_step_batch_size = batch_size
+            state.denoise_steps += 1
+            commit_length = int(commit_lengths[row])
+            chunk = commit_token_ids[row, :commit_length].detach().cpu().tolist()
+            state.last_delta_tokens = [int(token_id) for token_id in chunk]
+            before = len(state.generated_tokens)
+            try:
+                self.cache_pool.append(state, chunk)
+            except Exception as error:
+                self._finish(state, "error", error)
+                finished_ids.append(state.request_id)
+                continue
+            state.logical_length += commit_length
+            state.generated_tokens.extend(int(token_id) for token_id in chunk)
+            if self.continuous_batching_config.return_logprobs and commit_length:
+                probabilities = selected_confidence[row, :commit_length].clamp_min(
+                    torch.finfo(torch.float32).tiny
+                )
+                state.logprobs.extend(probabilities.log().detach().cpu().tolist())
+            if state.record_timestamps and commit_length:
+                state.timestamps.extend([time.perf_counter()] * commit_length)
+            if state.repetition_history is not None and commit_length:
+                tokens = commit_token_ids[row : row + 1]
+                eligible = commit_positions[row : row + 1]
+                _add_repetition_history(
+                    state.repetition_history.unsqueeze(0),
+                    tokens,
+                    eligible,
+                    self.excluded_repetition_token_ids,
+                )
+            state.jumps += int(jump_rows[row])
+            if bool(jump_rows[row]):
+                state.forced_jump_tokens += commit_length
+            if commit_length:
+                state.shifts += 1
+            state.rolling_state = shifted_states[row]
+
+            reason = None
+            if self.turn_end_token_id in chunk:
+                reason = "turn_end"
+            elif any(token_id in state.eos_token_ids for token_id in chunk):
+                reason = "eos"
+            elif len(state.generated_tokens) >= state.max_new_tokens:
+                reason = "max_new_tokens"
+            elif (
+                state.max_denoising_steps is not None
+                and state.denoise_steps >= state.max_denoising_steps
+            ):
+                reason = "max_denoising_steps"
+            elif state.denoise_steps >= state.max_iterations:
+                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)
+                finished_ids.append(state.request_id)
+            elif state.streaming and commit_length:
+                self._deliver(
+                    self._output_for(
+                        state,
+                        stream_update=True,
+                        delta_tokens=state.last_delta_tokens,
+                    )
+                )
+
+        self._stats["model_steps"] += 1
+        self._stats["generated_tokens"] += int(commit_lengths.sum())
+        self._stats["max_observed_batch_size"] = max(
+            self._stats["max_observed_batch_size"], batch_size
+        )
+        self._stats["active_slot_steps"] += batch_size
+        self._stats["slot_capacity_steps"] += self.max_requests_per_batch
+        return finished_ids
+
+    def _run_step_with_isolation(self, states: Sequence[ModilifyMk2RequestState]) -> None:
+        generator_states = {
+            state.request_id: state.generator.get_state()
+            for state in states
+            if state.generator is not None
+        }
+        try:
+            finished_ids = self._run_batch_step(states)
+        except Exception as batch_error:
+            for state in states:
+                if state.generator is not None:
+                    state.generator.set_state(generator_states[state.request_id])
+            if len(states) == 1:
+                self._finish(states[0], "error", batch_error)
+                finished_ids = [states[0].request_id]
+            else:
+                finished_ids = []
+                for state in states:
+                    if state.terminal_emitted:
+                        finished_ids.append(state.request_id)
+                        continue
+                    try:
+                        finished_ids.extend(self._run_batch_step([state]))
+                    except Exception as request_error:
+                        self._finish(state, "error", request_error)
+                        finished_ids.append(state.request_id)
+        with self._condition:
+            for request_id in dict.fromkeys(finished_ids):
+                state = self._active.pop(request_id, None)
+                if state is not None:
+                    self._active_reserved_blocks -= state.reserved_blocks
+            self._condition.notify_all()
+
+    @torch.inference_mode()
+    def _run_generation_loop(self) -> None:
+        try:
+            while True:
+                self._apply_cancellations()
+                if self._hard_stop:
+                    self._apply_cancellations()
+                with self._condition:
+                    has_active = bool(self._active)
+                # ``prefill_first`` fills every available slot before the next
+                # denoise step.  FIFO lets the already-active cohort take its
+                # next step first, then fills slots released by that step.
+                if (
+                    self.continuous_batching_config.scheduler_type == "prefill_first"
+                    or not has_active
+                ):
+                    self._admit_requests()
+                with self._condition:
+                    active = list(self._active.values())
+                    should_finish = (
+                        self._input_closed and not self._pending and not active
+                    )
+                    if should_finish:
+                        return
+                    if not active:
+                        self._condition.wait(timeout=0.05)
+                        continue
+                self._run_step_with_isolation(active)
+                if self.continuous_batching_config.scheduler_type == "fifo":
+                    self._admit_requests()
+        except BaseException as error:
+            self._fail_all_requests(error)
+        finally:
+            self._finished.set()
+            with self._condition:
+                self._condition.notify_all()
+
+
+@torch.inference_mode()
+def generate_static_batch_with_logical_cache(
+    model: Any,
+    input_ids: torch.LongTensor,
+    attention_mask: torch.BoolTensor | None,
+    generation_config: ModilifyMk2GenerationConfig,
+    *,
+    seeds: Sequence[int] | None = None,
+    max_new_tokens: Sequence[int] | None = None,
+) -> ModilifyMk2GenerationOutput:
+    """Run one fixed cohort through the same hole-free continuous engine."""
+
+    batch_size, input_width = input_ids.shape
+    if attention_mask is None:
+        attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+    else:
+        attention_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
+    if attention_mask.shape != input_ids.shape:
+        raise ValueError("`attention_mask` must have the same shape as `input_ids`.")
+    prompts = [
+        input_ids[row, attention_mask[row]].detach().cpu().tolist()
+        for row in range(batch_size)
+    ]
+    if any(not prompt for prompt in prompts):
+        raise ValueError("Every batched ModilifyMk2 prompt must contain at least one token.")
+    if seeds is not None and len(seeds) != batch_size:
+        raise ValueError("`seeds` must contain one seed per batch row.")
+    if max_new_tokens is None:
+        max_new_tokens = [int(generation_config.max_new_tokens)] * batch_size
+    if len(max_new_tokens) != batch_size or any(
+        not isinstance(limit, int) or isinstance(limit, bool) or limit <= 0
+        for limit in max_new_tokens
+    ):
+        raise ValueError("`max_new_tokens` must contain one positive limit per batch row.")
+
+    batching_config = ContinuousBatchingConfig(
+        block_size=max(4, int(getattr(model.config, "kv_cache_bucket_size", 128))),
+        max_batch_tokens=batch_size * int(model.config.canvas_length),
+        max_requests_per_batch=batch_size,
+        allow_block_sharing=False,
+        scheduler_type="prefill_first",
+    )
+    manager = ModilifyMk2ContinuousBatchingManager(
+        model=model,
+        generation_config=generation_config,
+        continuous_batching_config=batching_config,
+    )
+    try:
+        request_ids = []
+        for row, prompt in enumerate(prompts):
+            request_kwargs = {}
+            if seeds is not None:
+                request_kwargs["seed"] = int(seeds[row])
+            request_ids.append(
+                manager.add_request(
+                    prompt,
+                    request_id=f"static_{row}",
+                    max_new_tokens=int(max_new_tokens[row]),
+                    streaming=False,
+                    max_denoising_steps=generation_config.max_denoising_steps,
+                    eos_token_id=generation_config.eos_token_id,
+                    **request_kwargs,
+                )
+            )
+        manager.close_input()
+        # A static batch is one fixed cohort: queue every row before the worker
+        # starts so its first heavy forward necessarily contains the full batch.
+        manager.start()
+        final = {}
+        for output in manager:
+            if output.is_finished():
+                final[output.request_id] = output
+        ordered = [final[request_id] for request_id in request_ids]
+    finally:
+        manager.stop(block=True, hard_stop=True)
+        manager.destroy()
+    failures = [output for output in ordered if output.error is not None]
+    if failures:
+        details = "; ".join(
+            f"{output.request_id}: {output.error}" for output in failures
+        )
+        raise RuntimeError(f"Static ModilifyMk2 batch generation failed: {details}")
+
+    lengths = torch.tensor(
+        [len(output.generated_tokens) for output in ordered],
+        device=input_ids.device,
+        dtype=torch.long,
+    )
+    output_width = int(lengths.max()) if lengths.numel() else 0
+    pad_token_id = generation_config.pad_token_id
+    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)
+    generated = torch.full(
+        (batch_size, output_width),
+        pad_token_id,
+        device=input_ids.device,
+        dtype=input_ids.dtype,
+    )
+    for row, output in enumerate(ordered):
+        if output.generated_tokens:
+            generated[row, : len(output.generated_tokens)] = torch.tensor(
+                output.generated_tokens,
+                device=input_ids.device,
+                dtype=input_ids.dtype,
+            )
+
+    def tensor(name: str, *, dtype: torch.dtype) -> torch.Tensor:
+        return torch.tensor(
+            [getattr(output, name) for output in ordered],
+            device=input_ids.device,
+            dtype=dtype,
+        )
+
+    return ModilifyMk2GenerationOutput(
+        sequences=torch.cat((input_ids, generated), dim=-1),
+        generated_lengths=lengths,
+        tokens_per_forward=tensor("tokens_per_forward", dtype=torch.float32),
+        past_key_values=None,
+        stop_reason=tuple(output.stop_reason for output in ordered),
+        committed_tokens=lengths.clone(),
+        denoise_steps=tensor("denoise_steps", dtype=torch.long),
+        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",
+]

export_manifest.json

@@ -0,0 +1,103 @@
+{
+  "format": "modilify_mk2_transformers_inference_v1",
+  "global_step": 4400,
+  "state_schema_version": 25,
+  "training_scheme": "gold_prefix_shared_commit_0_256_committed_ce_calibration_causal_throughput_terminal_sft",
+  "memory_scheme": "dual_timescale_gdn2_trajectory_memory",
+  "memory_architecture": "compact_gdn2_v2",
+  "base_model": "google/diffusiongemma-26B-A4B-it",
+  "source_manifest_sha256": "e2b64fe2102ca1e002a62e279798abb3be271e3d9640a5afc5823cf233724dbc",
+  "precision_policy": "gdn2_small_fp32_v1",
+  "adapters": "unmerged",
+  "trainable_tensor_count": 632,
+  "total_tensor_count": 1323,
+  "shards": [
+    {
+      "file": "model-00001-of-00016.safetensors",
+      "bytes": 3898256004,
+      "sha256": "d3aa3fc0393a51acec27161569d566774030e0306a74748dd2a7a11fd8b535ec"
+    },
+    {
+      "file": "model-00002-of-00016.safetensors",
+      "bytes": 3798073636,
+      "sha256": "89339fd6f0521977fff9d4ae498e485ff409dbb8716cc924c8667c8bceaa19cd"
+    },
+    {
+      "file": "model-00003-of-00016.safetensors",
+      "bytes": 3290562732,
+      "sha256": "b9351e59c7c604b26712960b5d9f3460f5c1aace9ccb348999a2334d93fb613d"
+    },
+    {
+      "file": "model-00004-of-00016.safetensors",
+      "bytes": 3319472948,
+      "sha256": "1c89ba68a1f56e986dbf43175a7a7a6f2d046b04ac993c2837024f71b5d5923d"
+    },
+    {
+      "file": "model-00005-of-00016.safetensors",
+      "bytes": 3290562732,
+      "sha256": "4e42f45a5c409649c0ac8a262b26a0f644d10619d5f857568db855d12a865183"
+    },
+    {
+      "file": "model-00006-of-00016.safetensors",
+      "bytes": 3290562772,
+      "sha256": "da6e7d99f9801cb1b7007044f3baa66e767412d4350211a637174321eacee401"
+    },
+    {
+      "file": "model-00007-of-00016.safetensors",
+      "bytes": 3319473028,
+      "sha256": "ce1f462e324e162f9c0f41b1db3945290eb06c520da3d8ff0f582cde79a6fa96"
+    },
+    {
+      "file": "model-00008-of-00016.safetensors",
+      "bytes": 3290562812,
+      "sha256": "a163f6df13b9a3fbe1df187e42700b74d71fa08b331c521be568f71dec77a8f0"
+    },
+    {
+      "file": "model-00009-of-00016.safetensors",
+      "bytes": 3290562812,
+      "sha256": "eacc2e9ac62fa45f4fd1d391bfff490e327c19df3f70e7835e556287fc9926bd"
+    },
+    {
+      "file": "model-00010-of-00016.safetensors",
+      "bytes": 3319473028,
+      "sha256": "5fa9c8080a9c3954985226496dee261938f3d05cf81e4436932de593e5a01ae1"
+    },
+    {
+      "file": "model-00011-of-00016.safetensors",
+      "bytes": 3290562812,
+      "sha256": "7df1df1cea73133a94aec88056e2514e04a3467d42c0dd55ad0b922af2e5e986"
+    },
+    {
+      "file": "model-00012-of-00016.safetensors",
+      "bytes": 3290562812,
+      "sha256": "797078b4703d602cc3a323dba676fa38bb959b70d455bdc8c673663261eee2d4"
+    },
+    {
+      "file": "model-00013-of-00016.safetensors",
+      "bytes": 3319473028,
+      "sha256": "35923e8f76bc78ae89b3f83d1cdd175649f64cc898f1862baf543d7cdfbcd0e1"
+    },
+    {
+      "file": "model-00014-of-00016.safetensors",
+      "bytes": 3290562812,
+      "sha256": "105ceea4ed17c811fb79036c199db1f7d2392a5cc4c05003085f7217e8a938bb"
+    },
+    {
+      "file": "model-00015-of-00016.safetensors",
+      "bytes": 3290562812,
+      "sha256": "d3569ec55bc4f7e89de00efaedfbca3df44096526291731ec14bd7fc245c3b39"
+    },
+    {
+      "file": "model-00016-of-00016.safetensors",
+      "bytes": 1202380654,
+      "sha256": "23c9c00fc04b475c744458f498178a6e9a20fd673c36af981dd22e963b7ca49b"
+    }
+  ],
+  "validated_environment": {
+    "python": "3.12",
+    "torch": "2.14.0",
+    "transformers": "5.14.1",
+    "safetensors": "0.8.0",
+    "accelerate": "1.14.0"
+  }
+}

gdn2_memory.py

@@ -0,0 +1,119 @@
+"""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 empty(self, *leading: int, device: torch.device | None = None) -> torch.Tensor:
+        return torch.zeros(*leading, self.heads, self.key_dim, self.value_dim,
+                           device=device, dtype=torch.float32)
+
+    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
+
+
+__all__ = ["GDN2Memory"]

gdn2_trajectory.py

@@ -0,0 +1,121 @@
+"""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 detach(self) -> "GDN2TrajectoryState":
+        return GDN2TrajectoryState(
+            self.cells.detach(), self.row.detach(), self.persistent.detach(),
+            self.seen.detach(),
+        )
+
+    def clear_refill(self, lengths: torch.Tensor, head: torch.Tensor) -> "GDN2TrajectoryState":
+        batch, canvas = self.seen.shape
+        if lengths.shape != (batch,) or head.shape != (batch,):
+            raise ValueError("Commit lengths and canvas heads must be per row.")
+        physical = torch.arange(canvas, device=self.seen.device)[None, :]
+        overwritten = ((physical - head[:, None]) % canvas) < lengths[:, None]
+        return GDN2TrajectoryState(
+            self.cells.masked_fill(overwritten[..., None, None, None], 0.0),
+            self.row, self.persistent, self.seen & ~overwritten,
+        )
+
+    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 empty(self, batch: int, canvas: int, *, device: torch.device | None = None
+              ) -> GDN2TrajectoryState:
+        return GDN2TrajectoryState(
+            self.cell.empty(batch, canvas, device=device),
+            self.row.empty(batch, device=device),
+            self.persistent.empty(batch, device=device),
+            torch.zeros(batch, canvas, device=device, dtype=torch.bool),
+        )
+
+    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)
+
+    def commit(self, state: GDN2TrajectoryState, experience: torch.Tensor,
+               valid: torch.Tensor) -> GDN2TrajectoryState:
+        if experience.shape[:2] != valid.shape:
+            raise ValueError("Commit experience and validity shapes differ.")
+        long_state = self.persistent.write_sequence(state.persistent, experience, valid)
+        return GDN2TrajectoryState(state.cells, state.row, long_state, state.seen)
+
+
+__all__ = ["GDN2TrajectoryMemory", "GDN2TrajectoryState"]

generation_config.json

@@ -0,0 +1,48 @@
+{
+  "bos_token_id": 2,
+  "confidence_threshold": null,
+  "eos_token_id": [
+    1,
+    50
+  ],
+  "jump_on_no_progress_after": 12,
+  "max_denoising_steps": 256,
+  "max_new_tokens": 256,
+  "max_ponder_steps": 64,
+  "min_trajectory_progress": 0.005,
+  "pad_token_id": 0,
+  "repetition_penalty": 1.0,
+  "repetition_penalty_exclude_token_ids": [
+    2,
+    1,
+    3,
+    0,
+    4,
+    258881,
+    256000,
+    255999,
+    258883,
+    101,
+    258882,
+    106,
+    52,
+    49,
+    47,
+    51,
+    258880,
+    100,
+    105,
+    48,
+    46,
+    50,
+    98,
+    258884
+  ],
+  "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

@@ -0,0 +1,1307 @@
+"""Rolling text generation for latent-memory ModilifyMk2."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, replace
+import math
+import time
+from typing import Any
+
+import torch
+from transformers.cache_utils import Cache
+from transformers.generation import LogitsProcessorList
+from transformers.generation.streamers import BaseStreamer
+from transformers.modeling_outputs import ModelOutput
+
+from transformers.models.diffusion_gemma import (
+    DiffusionGemmaGenerationConfig,
+    DiffusionGemmaGenerationMixin,
+)
+from .commit_policy import (
+    fused_commit_failure_rate,
+    select_commit_lengths,
+)
+from .configuration_modilify_mk2 import DENOISE_TEMPERATURE
+from .latent_deliberation import LatentDeliberationState, infer_commit_reason
+
+
+_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",
+    "bos_token_id",
+    "pad_token_id",
+    "eos_token_id",
+    "_commit_hash",
+    "_from_model_config",
+    "transformers_version",
+})
+_UNSUPPORTED_GENERATION_FIELDS = (
+    "sampler_config",
+    "stability_threshold",
+    "confidence_threshold",
+    "one_token_per_denoise_step",
+    "compile_generation",
+    "sliding_denoise",
+    "adaptive_ponder_budget",
+    "force_commit_on_max_steps",
+    "ponder_budget_id",
+)
+
+
+def _reject_unsupported_generation_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:
+            continue
+        if isinstance(value, int):
+            token_ids.add(int(value))
+            continue
+        if isinstance(value, (list, tuple, set)):
+            token_ids.update(int(token_id) for token_id in value if token_id is not None)
+    return token_ids
+
+
+def _add_repetition_history(
+    history: torch.BoolTensor,
+    token_ids: torch.LongTensor,
+    eligible: torch.BoolTensor,
+    excluded_token_ids: set[int],
+) -> None:
+    """Add eligible row-local token IDs to a compact [batch, vocab] history."""
+
+    if token_ids.shape != eligible.shape or token_ids.shape[0] != history.shape[0]:
+        raise ValueError("Repetition history token and eligibility shapes must match.")
+    eligible = eligible.clone()
+    for token_id in excluded_token_ids:
+        eligible &= token_ids.ne(token_id)
+    if not bool(eligible.any()):
+        return
+    rows = torch.arange(history.shape[0], device=history.device)[:, None]
+    rows = rows.expand_as(token_ids)
+    history[rows[eligible], token_ids[eligible]] = True
+
+
+class ModilifyMk2GenerationConfig(DiffusionGemmaGenerationConfig):
+    def __init__(self, **kwargs):
+        _reject_unsupported_generation_fields({
+            name: kwargs.pop(name)
+            for name in _UNSUPPORTED_GENERATION_FIELDS
+            if name in kwargs
+        })
+        self.turn_end_token_id: int | None = kwargs.pop("turn_end_token_id", None)
+        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))
+        self.repetition_penalty: float = float(kwargs.pop("repetition_penalty", 1.0))
+        excluded_token_ids = kwargs.pop("repetition_penalty_exclude_token_ids", ())
+        self.repetition_penalty_exclude_token_ids: list[int] = list(
+            dict.fromkeys(int(token_id) for token_id in excluded_token_ids or ())
+        )
+        kwargs.pop("t_min", None)
+        kwargs.pop("t_max", None)
+        parent_kwargs = {
+            name: kwargs.pop(name)
+            for name in tuple(kwargs)
+            if name in _PARENT_GENERATION_KEYS
+        }
+        super().__init__(
+            sampler_config=None,
+            stability_threshold=None,
+            confidence_threshold=None,
+            **parent_kwargs,
+        )
+        self.t_min = DENOISE_TEMPERATURE
+        self.t_max = DENOISE_TEMPERATURE
+        self.validate()
+
+    def update(self, defaults_only=False, allow_custom_entries=False, **kwargs):
+        """Apply supported overrides while keeping temperature globally fixed."""
+
+        _reject_unsupported_generation_fields({
+            name: kwargs.pop(name)
+            for name in _UNSUPPORTED_GENERATION_FIELDS
+            if name in kwargs
+        })
+        if "turn_end_token_id" in kwargs:
+            self.turn_end_token_id = kwargs.pop("turn_end_token_id")
+        if "max_ponder_steps" in kwargs:
+            self.max_ponder_steps = kwargs.pop("max_ponder_steps")
+        if "jump_on_no_progress_after" in kwargs:
+            self.jump_on_no_progress_after = kwargs.pop("jump_on_no_progress_after")
+        if "min_trajectory_progress" in kwargs:
+            self.min_trajectory_progress = float(kwargs.pop("min_trajectory_progress"))
+        if "repetition_penalty" in kwargs:
+            self.repetition_penalty = float(kwargs.pop("repetition_penalty"))
+        if "repetition_penalty_exclude_token_ids" in kwargs:
+            excluded_token_ids = kwargs.pop("repetition_penalty_exclude_token_ids")
+            self.repetition_penalty_exclude_token_ids = list(
+                dict.fromkeys(int(token_id) for token_id in excluded_token_ids or ())
+            )
+        kwargs.pop("t_min", None)
+        kwargs.pop("t_max", None)
+        unused = super().update(
+            defaults_only=defaults_only,
+            allow_custom_entries=allow_custom_entries,
+            **kwargs,
+        )
+        self.sampler_config = None
+        self.stability_threshold = None
+        self.confidence_threshold = None
+        self.t_min = DENOISE_TEMPERATURE
+        self.t_max = DENOISE_TEMPERATURE
+        return unused
+
+    def validate(self, **kwargs):
+        if self.max_new_tokens is not None and (
+            not isinstance(self.max_new_tokens, int) or self.max_new_tokens <= 0
+        ):
+            raise ValueError(f"`max_new_tokens` must be a positive integer, but got {self.max_new_tokens}")
+        if self.max_length is not None and (
+            not isinstance(self.max_length, int) or self.max_length <= 0
+        ):
+            raise ValueError(f"`max_length` must be a positive integer, but got {self.max_length}")
+        if self.turn_end_token_id is not None and (
+            not isinstance(self.turn_end_token_id, int) or self.turn_end_token_id < 0
+        ):
+            raise ValueError("`turn_end_token_id` must be a non-negative integer.")
+        if not isinstance(self.max_ponder_steps, int) or self.max_ponder_steps <= 0:
+            raise ValueError("`max_ponder_steps` must be a positive integer.")
+        if not isinstance(self.jump_on_no_progress_after, int) or self.jump_on_no_progress_after <= 0:
+            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.repetition_penalty) or self.repetition_penalty <= 0:
+            raise ValueError("`repetition_penalty` must be a finite positive number.")
+        if any(
+            not isinstance(token_id, int) or isinstance(token_id, bool) or token_id < 0
+            for token_id in self.repetition_penalty_exclude_token_ids
+        ):
+            raise ValueError(
+                "`repetition_penalty_exclude_token_ids` must contain non-negative integers."
+            )
+
+    @classmethod
+    def from_model_config(cls, model_config):
+        """Build the only generation field owned by the ModilifyMk2 model config."""
+
+        return cls(turn_end_token_id=model_config.turn_end_token_id)
+
+    @staticmethod
+    def _get_default_generation_params() -> dict[str, object]:
+        """Return defaults with no inherited entropy/readiness commit controls."""
+
+        return {
+            "max_new_tokens": 256,
+            "max_denoising_steps": 48,
+            "t_min": DENOISE_TEMPERATURE,
+            "t_max": DENOISE_TEMPERATURE,
+        }
+
+
+@dataclass
+class ModilifyMk2GenerationOutput(ModelOutput):
+    sequences: torch.LongTensor
+    generated_lengths: torch.LongTensor | None = None
+    tokens_per_forward: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    stop_reason: str | tuple[str, ...] | None = None
+    committed_tokens: int | torch.LongTensor | None = None
+    denoise_steps: int | torch.LongTensor | None = None
+    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
+    head: torch.LongTensor | None = None
+
+
+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().flatten(2)
+    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.latent_state.gdn2.seen[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,
+    }
+
+
+class NoiseCanvasSampler:
+    """Uniform diffusion noise source with no commit-policy responsibilities."""
+
+    def __init__(self, *, canvas_length: int, vocab_size: int) -> None:
+        self.canvas_length = int(canvas_length)
+        self.vocab_size = int(vocab_size)
+        self.initial_entropy = math.log(self.vocab_size)
+
+    def initialize_canvas(
+        self,
+        batch_size: int,
+        device: torch.device,
+        generators: Sequence[torch.Generator] | None = None,
+    ) -> torch.LongTensor:
+        if generators is not None:
+            if len(generators) != batch_size:
+                raise ValueError("Canvas sampling requires one generator per batch row.")
+            return torch.cat(
+                [
+                    torch.randint(
+                        self.vocab_size,
+                        (1, self.canvas_length),
+                        device=device,
+                        generator=generator,
+                    )
+                    for generator in generators
+                ],
+                dim=0,
+            )
+        return torch.randint(
+            self.vocab_size,
+            (batch_size, self.canvas_length),
+            device=device,
+        )
+
+
+class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
+    def init_continuous_batching(
+        self,
+        generation_config=None,
+        continuous_batching_config=None,
+        workload_hints=None,
+    ):
+        """Create the ModilifyMk2 continuous manager behind the official API surface."""
+
+        from .continuous_batching import (
+            ModilifyMk2ContinuousBatchingManager,
+            continuous_config_fingerprint,
+        )
+
+        cached = getattr(self, "_cached_continuous_batching_manager", None)
+        if isinstance(cached, ModilifyMk2ContinuousBatchingManager) and not cached.destroyed:
+            requested_generation = generation_config or getattr(
+                self, "generation_config", None
+            )
+            requested_fingerprint = continuous_config_fingerprint(
+                requested_generation,
+                continuous_batching_config,
+            )
+            if cached.config_fingerprint == requested_fingerprint:
+                cached._prepare_for_next_session()
+                return cached
+            cached.destroy()
+            delattr(self, "_cached_continuous_batching_manager")
+        return ModilifyMk2ContinuousBatchingManager(
+            model=self,
+            generation_config=generation_config or getattr(self, "generation_config", None),
+            continuous_batching_config=continuous_batching_config,
+            workload_hints=workload_hints,
+        )
+
+    def destroy_cached_continuous_batching_manager(self) -> None:
+        manager = getattr(self, "_cached_continuous_batching_manager", None)
+        if manager is not None:
+            manager.destroy()
+            delattr(self, "_cached_continuous_batching_manager")
+
+    @contextmanager
+    @torch.no_grad()
+    def continuous_batching_context_manager(
+        self,
+        generation_config=None,
+        block: bool = True,
+        timeout: float | None = None,
+        continuous_batching_config=None,
+        persistent_manager: bool = False,
+        warmup: bool = True,
+        workload_hints=None,
+    ):
+        manager = self.init_continuous_batching(
+            generation_config=generation_config,
+            continuous_batching_config=continuous_batching_config,
+            workload_hints=workload_hints,
+        )
+        if persistent_manager:
+            self._cached_continuous_batching_manager = manager
+        if warmup and not manager.warmed_up:
+            manager.warmup()
+        manager.start()
+        try:
+            yield manager
+        finally:
+            manager.stop(
+                block=block,
+                timeout=timeout,
+                keep_for_next_session=persistent_manager,
+            )
+            if not persistent_manager:
+                manager.destroy()
+
+    @torch.no_grad()
+    def generate_batch(
+        self,
+        inputs: list[list[int]],
+        generation_config=None,
+        continuous_batching_config=None,
+        record_timestamps: bool = False,
+        progress_bar: bool = True,
+        persistent_manager: bool = False,
+        warmup: bool = True,
+        **kwargs: Any,
+    ) -> dict[str, object]:
+        """Official-compatible convenience wrapper over the continuous manager."""
+
+        del progress_bar
+        if any(
+            not isinstance(input_ids, list)
+            or not input_ids
+            or any(
+                not isinstance(token_id, int) or isinstance(token_id, bool)
+                for token_id in input_ids
+            )
+            for input_ids in inputs
+        ):
+            raise ValueError("Every `inputs` row must be a non-empty list of integer token IDs.")
+        seeds = kwargs.pop("seeds", None)
+        if seeds is not None and len(seeds) != len(inputs):
+            raise ValueError("`seeds` must contain one seed per request.")
+        if seeds is not None:
+            seeds = [int(seed) for seed in seeds]
+        if not inputs:
+            return {}
+        manager = self.init_continuous_batching(
+            generation_config=generation_config,
+            continuous_batching_config=continuous_batching_config,
+        )
+        if persistent_manager:
+            self._cached_continuous_batching_manager = manager
+        if warmup and not manager.warmed_up:
+            manager.warmup()
+        completed = False
+        try:
+            request_ids = []
+            queue_limit = int(
+                manager.continuous_batching_config.max_queue_size or 0
+            )
+            preload_count = (
+                len(inputs) if queue_limit == 0 else min(len(inputs), queue_limit)
+            )
+            for index, input_ids in enumerate(inputs):
+                if index == preload_count and not manager.is_running():
+                    manager.start()
+                request_kwargs = dict(kwargs)
+                if seeds is not None:
+                    request_kwargs["seed"] = int(seeds[index])
+                request_ids.append(
+                    manager.add_request(
+                        input_ids=input_ids,
+                        record_timestamps=record_timestamps,
+                        **request_kwargs,
+                    )
+                )
+            manager.close_input()
+            # Unlike the open-ended manager API, generate_batch has its whole
+            # initial workload. Queue it before starting so the first cohort is
+            # filled deterministically up to scheduler/queue capacity.
+            if not manager.is_running():
+                manager.start()
+            final_outputs = {}
+            for output in manager:
+                if output.status in {
+                    getattr(output.status.__class__, "FINISHED", output.status),
+                    getattr(output.status.__class__, "FAILED", output.status),
+                }:
+                    final_outputs[output.request_id] = output
+            result = {
+                request_id: final_outputs[request_id]
+                for request_id in request_ids
+                if request_id is not None
+            }
+            completed = True
+        finally:
+            manager.stop(
+                block=True,
+                keep_for_next_session=persistent_manager,
+                hard_stop=not completed,
+            )
+            if not persistent_manager:
+                manager.destroy()
+        return result
+
+    def _prepare_sampler(
+        self, generation_config: ModilifyMk2GenerationConfig, canvas_length: int | None = None
+    ) -> NoiseCanvasSampler:
+        del generation_config
+        return NoiseCanvasSampler(
+            canvas_length=canvas_length or self.config.canvas_length,
+            vocab_size=self.config.text_config.vocab_size,
+        )
+
+    def _write_committed_memory(
+        self,
+        *,
+
+        next_state: ModilifyMk2RollingState,
+        working_state: torch.Tensor,
+
+        heavy_hidden: torch.Tensor,
+        commit_token_ids: torch.LongTensor,
+        commit_lengths: torch.Tensor,
+        prefix_lengths: torch.Tensor,
+        commit_reason: torch.Tensor,
+        canvas_head: torch.Tensor | None = None,
+    ) -> ModilifyMk2RollingState:
+        if not bool(commit_lengths.gt(0).any()):
+            return next_state
+        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, _diagnostics = self.latent_deliberation.commit_write(
+            memory=next_state.latent_state.memory_slots,
+            working_state=working_state,
+
+
+            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,
+                gdn2=replace(next_state.latent_state.gdn2, persistent=memory),
+            ),
+        )
+
+    @staticmethod
+    def _shift_state_rows(
+        state: ModilifyMk2RollingState,
+        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."""
+
+        batch_size, canvas_length = state.canvas.shape
+        if commit_lengths.shape != (batch_size,):
+            raise ValueError("Commit lengths must have shape [batch].")
+        if not bool(commit_lengths.gt(0).any()):
+            if generators is not None:
+                if len(generators) != batch_size:
+                    raise ValueError("State shifting requires one generator per batch row.")
+                # Seeded generation deliberately advances every active request
+                # once per denoise step, independent of the other active rows.
+                for generator in generators:
+                    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]
+        retained = source.lt(canvas_length)
+
+        def shift(value: torch.Tensor, fill_value: float | int = 0) -> torch.Tensor:
+            index = source.clamp_max(canvas_length - 1)
+            index = index.view(
+                batch_size, canvas_length, *([1] * (value.ndim - 2))
+            ).expand_as(value)
+            gathered = value.gather(1, index)
+            mask = retained.view(
+                batch_size, canvas_length, *([1] * (value.ndim - 2))
+            )
+            fill = torch.as_tensor(fill_value, device=value.device, dtype=value.dtype)
+            return torch.where(mask, gathered, fill)
+
+        if generators is None:
+            tail = sampler.initialize_canvas(batch_size, state.canvas.device)
+        else:
+            if len(generators) != batch_size:
+                raise ValueError("State shifting requires one generator per batch row.")
+            tail = torch.zeros_like(state.canvas)
+            for row, (commit_length, generator) in enumerate(
+                zip(commit_lengths.detach().cpu().tolist(), generators, strict=True)
+            ):
+                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)
+        shifted_latent = LatentDeliberationState(
+            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=None if latent.gdn2 is None else 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,
+
+
+            head=state.head,
+        )
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        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)
+        scalar_seed = kwargs.pop("seed", None)
+        if request_seeds is not None and scalar_seed is not None:
+            raise ValueError("Pass either `seed` or `seeds`, not both.")
+        generation_config, model_kwargs = self._prepare_generation_config(generation_config, **kwargs)
+        if input_ids is None or input_ids.ndim != 2 or input_ids.shape[0] < 1:
+            raise ValueError("ModilifyMk2 generation requires `input_ids` with shape [batch, sequence].")
+        if logits_processor:
+            raise ValueError(
+                "ModilifyMk2 uses its built-in exact sampler and does not accept "
+                "custom logits processors."
+            )
+        batch_size, input_width = input_ids.shape
+        if scalar_seed is not None:
+            request_seeds = [int(scalar_seed) + row for row in range(batch_size)]
+        elif batch_size > 1 and request_seeds is None:
+            # Static B>1 still uses independent row generators, but its default
+            # path must advance the caller's global device RNG just like normal
+            # generation instead of repeating torch.initial_seed() forever.
+            seed_parts = torch.randint(
+                0,
+                (1 << 31) - 1,
+                (batch_size, 2),
+                device=input_ids.device,
+                dtype=torch.int64,
+            ).detach().cpu().tolist()
+            request_seeds = [
+                (int(high) << 31) | int(low) for high, low in seed_parts
+            ]
+        if request_seeds is not None:
+            if len(request_seeds) != batch_size:
+                raise ValueError("`seeds` must contain one seed per batch row.")
+            sampling_generators = []
+            for seed_value in request_seeds:
+                generator = torch.Generator(device=input_ids.device)
+                generator.manual_seed(int(seed_value) & ((1 << 63) - 1))
+                sampling_generators.append(generator)
+        else:
+            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:
+            from .continuous_batching import generate_static_batch_with_logical_cache
+
+            attention_mask = model_kwargs.pop("attention_mask", None)
+            canonical_mask = (
+                torch.ones_like(input_ids, dtype=torch.bool)
+                if attention_mask is None
+                else attention_mask.to(device=input_ids.device, dtype=torch.bool)
+            )
+            if canonical_mask.shape != input_ids.shape:
+                raise ValueError(
+                    "`attention_mask` must have the same shape as `input_ids`."
+                )
+            row_limits = []
+            for prompt_length in canonical_mask.long().sum(dim=-1).tolist():
+                _, resolved_max_new_tokens = self._prepare_generated_length(
+                    generation_config, int(prompt_length)
+                )
+                if resolved_max_new_tokens <= 0:
+                    raise ValueError(
+                        "The requested maximum length leaves no room to generate "
+                        "for every batch row."
+                    )
+                row_limits.append(int(resolved_max_new_tokens))
+            provided_positions = model_kwargs.pop("position_ids", None)
+            if provided_positions is not None:
+                canonical_positions = (
+                    canonical_mask.long().cumsum(dim=-1).sub(1).clamp_min(0)
+                ).to(provided_positions)
+                if not torch.equal(provided_positions, canonical_positions):
+                    raise ValueError(
+                        "Continuous static batches require canonical row-local `position_ids`."
+                    )
+            if model_kwargs:
+                unsupported = ", ".join(sorted(model_kwargs))
+                raise ValueError(
+                    f"Unsupported batched ModilifyMk2 generation arguments: {unsupported}"
+                )
+            return generate_static_batch_with_logical_cache(
+                self,
+                input_ids,
+                canonical_mask,
+                generation_config,
+                seeds=request_seeds,
+                max_new_tokens=row_limits,
+            )
+        device = input_ids.device
+        dtype = self.model.decoder.embed_tokens.weight.dtype
+        canvas_length = self.config.canvas_length
+        cached_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        repetition_penalty = float(generation_config.repetition_penalty)
+        repetition_enabled = repetition_penalty != 1.0
+        if repetition_enabled and cached_length:
+            raise ValueError(
+                "Repetition penalty requires a fresh KV cache so the complete prompt "
+                "token history is available."
+            )
+        _, 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,
+        )
+        if past_key_values is None:
+            past_key_values = self._prepare_cache_for_generation(
+                generation_config,
+                batch_size=batch_size,
+                # Ragged batches append dense, masked cache blocks. In the
+                # worst case only one row advances in each block.
+                max_length=input_width + batch_size * max_new_tokens,
+            )
+        expected_mask_width = cached_length + input_width
+        cache_attention_mask = model_kwargs.pop(
+            "attention_mask",
+            torch.ones(
+                batch_size, expected_mask_width, dtype=torch.bool, device=device
+            ),
+        ).bool()
+        if cache_attention_mask.shape != (batch_size, expected_mask_width):
+            raise ValueError(
+                "`attention_mask` must have shape [batch, cached_length + sequence]."
+            )
+        provided_position_ids = model_kwargs.pop("position_ids", None)
+        if provided_position_ids is not None:
+            if provided_position_ids.shape != input_ids.shape:
+                raise ValueError("`position_ids` must have the same shape as `input_ids`.")
+            prompt_positions = provided_position_ids.to(device=device, dtype=torch.int32)
+        elif cached_length:
+            prompt_positions = torch.arange(
+                cached_length,
+                cached_length + input_width,
+                device=device,
+                dtype=torch.int32,
+            ).unsqueeze(0)
+        else:
+            input_mask = cache_attention_mask[:, -input_width:]
+            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:
+            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,
+            ).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, enable_gdn2=True,
+        )
+        initial_canvas = sampler.initialize_canvas(
+            batch_size, device, generators=sampling_generators
+        )
+        state = ModilifyMk2RollingState(
+            canvas=initial_canvas,
+            confidence=torch.zeros(
+                batch_size, canvas_length, device=device, dtype=torch.float32
+            ),
+            entropy=torch.full(
+                (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,
+
+
+            head=torch.zeros(batch_size, device=device, dtype=torch.long),
+        )
+        turn_end = (
+            self.config.turn_end_token_id
+            if generation_config.turn_end_token_id is None
+            else generation_config.turn_end_token_id
+        )
+        configured_eos = generation_config.eos_token_id
+        if configured_eos is None:
+            configured_eos = self.config.eos_token_id
+        if isinstance(configured_eos, int):
+            configured_eos = [configured_eos]
+        stop_token_ids = tuple(
+            dict.fromkeys((int(turn_end), *(int(value) for value in configured_eos or ())))
+        )
+        pad_token_id = generation_config.pad_token_id
+        if pad_token_id is None:
+            pad_token_id = getattr(self.config, "pad_token_id", None)
+        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,
+            generation_config.bos_token_id,
+            generation_config.eos_token_id,
+            generation_config.turn_end_token_id,
+            getattr(self.config, "image_token_id", None),
+        )
+        repetition_history = None
+        if repetition_enabled:
+            repetition_history = torch.zeros(
+                (batch_size, self.config.text_config.vocab_size),
+                dtype=torch.bool,
+                device=device,
+            )
+            _add_repetition_history(
+                repetition_history,
+                input_ids,
+                cache_attention_mask[:, -input_width:],
+                excluded_repetition_token_ids,
+            )
+        generated = torch.full(
+            (batch_size, max_new_tokens),
+            pad_token_id,
+            dtype=input_ids.dtype,
+            device=device,
+        )
+        committed = torch.zeros(batch_size, dtype=torch.long, device=device)
+        denoise_steps = torch.zeros_like(committed)
+        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, :]
+        if streamer is not None:
+            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, :]
+            ).to(torch.int32)
+            denoise_steps += active_rows.long()
+            decoder_attention_mask = torch.cat(
+                (
+                    cache_attention_mask,
+                    torch.ones(
+                        batch_size,
+                        canvas_length,
+                        dtype=torch.bool,
+                        device=device,
+                    ),
+                ),
+                dim=-1,
+            )
+            output = self(
+                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,
+
+
+                decoder_position_ids=decoder_positions, decoder_read_cache=True,
+                decoder_attention_mask=decoder_attention_mask,
+                compact_vocab=True,
+                repetition_token_mask=repetition_history,
+                repetition_penalty=repetition_penalty,
+                sampling_generators=sampling_generators,
+                **model_kwargs,
+            )
+            if (
+                output.proposal is None
+                or output.proposal_confidence is None
+                or output.token_entropy is None
+                or output.greedy_proposal is None
+                or output.greedy_confidence is None
+            ):
+                raise RuntimeError(
+                    "Compact vocabulary forward did not return proposal statistics."
+                )
+            proposal = output.proposal
+            proposal_confidence = output.proposal_confidence
+            token_entropy = output.token_entropy
+            greedy_proposal = output.greedy_proposal
+            greedy_confidence = output.greedy_confidence
+            next_canvas = proposal.clone()
+            next_confidence = proposal_confidence.float()
+            next_latent = replace(
+                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)
+            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,
+                latent_state=next_latent,
+
+
+                head=state.head,
+            )
+            normal_failure_rate = fused_commit_failure_rate(
+                proposal_confidence, token_entropy,
+                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),
+                vocab_size=self.config.text_config.vocab_size,
+            )
+            jump_failure_rate = fused_commit_failure_rate(
+                greedy_confidence, token_entropy,
+                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),
+                vocab_size=self.config.text_config.vocab_size,
+            )
+            previous_failure_rate = fused_commit_failure_rate(
+                state.confidence, state.entropy,
+                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),
+                vocab_size=self.config.text_config.vocab_size,
+            )
+            policy_decision = select_commit_lengths(
+                sampled_token_ids=proposal,
+                normal_failure_rate=normal_failure_rate,
+                previous_failure_rate=previous_failure_rate,
+                greedy_token_ids=greedy_proposal,
+                jump_failure_rate=jump_failure_rate,
+                ponder_steps=state.latent_state.ponder_steps,
+                stagnation_steps=state.latent_state.stagnation_steps,
+                active_rows=active_rows,
+                remaining_lengths=remaining,
+                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
+            jump_rows = policy_decision.jump_rows
+            jumps += jump_rows.long()
+            forced_jump_tokens += torch.where(
+                jump_rows, commit_lengths, torch.zeros_like(commit_lengths)
+            )
+            commit_positions = canvas_positions.lt(commit_lengths[:, None])
+            if bool(jump_rows.any()):
+                next_state = replace(
+                    next_state,
+                    canvas=torch.where(
+                        commit_positions & jump_rows[:, None],
+                        policy_decision.commit_token_ids,
+                        next_state.canvas,
+                    ),
+                )
+            next_state = replace(
+                next_state,
+                latent_state=replace(
+                    next_state.latent_state,
+                    ponder_steps=next_ponder,
+                    stagnation_steps=next_stagnation,
+                ),
+            )
+            commit_token_ids = policy_decision.commit_token_ids
+            before = committed.clone()
+            write_rows = torch.arange(batch_size, device=device)[:, None].expand_as(
+                commit_token_ids
+            )
+            write_positions = before[:, None] + canvas_positions
+            generated[
+                write_rows[commit_positions], write_positions[commit_positions]
+            ] = commit_token_ids[commit_positions]
+            if repetition_history is not None:
+                _add_repetition_history(
+                    repetition_history,
+                    commit_token_ids,
+                    commit_positions,
+                    excluded_repetition_token_ids,
+                )
+
+            commit_width = int(commit_lengths.max())
+            if commit_width:
+                block_mask = torch.arange(commit_width, device=device)[None, :].lt(
+                    commit_lengths[:, None]
+                )
+                committed_block = torch.where(
+                    block_mask,
+                    commit_token_ids[:, :commit_width],
+                    torch.full(
+                        (batch_size, commit_width),
+                        pad_token_id,
+                        device=device,
+                        dtype=input_ids.dtype,
+                    ),
+                )
+                block_positions = (
+                    logical_lengths[:, None]
+                    + torch.arange(commit_width, device=device)[None, :]
+                ).to(torch.int32)
+                block_positions = torch.where(
+                    block_mask, block_positions, torch.zeros_like(block_positions)
+                )
+                cache_attention_mask = torch.cat(
+                    (cache_attention_mask, block_mask), dim=-1
+                )
+                past_key_values = self.model.encoder(
+                    input_ids=committed_block,
+                    attention_mask=cache_attention_mask,
+                    past_key_values=past_key_values,
+                    position_ids=block_positions,
+                ).past_key_values
+                if streamer is not None:
+                    streamer.put(committed_block.cpu())
+            committed_rows = commit_lengths.gt(0)
+            shifts += committed_rows.long()
+            if (
+                output.working_state is None
+            ):
+                raise RuntimeError("Forward did not return working trajectory features.")
+            next_state = self._write_committed_memory(
+
+                next_state=next_state,
+                working_state=output.working_state,
+
+                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(
+                    commit_lengths,
+                    jump_rows=jump_rows,
+                    commit_token_ids=commit_token_ids,
+                    terminal_token_ids=stop_token_ids,
+                ),
+            )
+            shifted = self._shift_state_rows(
+                next_state,
+                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
+            ).any(dim=-1)
+            eos_hits = torch.zeros_like(turn_hits)
+            for token_id in stop_token_ids:
+                if token_id != turn_end:
+                    eos_hits |= (
+                        commit_token_ids.eq(token_id) & commit_positions
+                    ).any(dim=-1)
+            stop_codes = torch.where(
+                stop_codes.eq(0) & turn_hits,
+                torch.ones_like(stop_codes),
+                stop_codes,
+            )
+            stop_codes = torch.where(
+                stop_codes.eq(0) & eos_hits,
+                torch.full_like(stop_codes, 2),
+                stop_codes,
+            )
+            stop_codes = torch.where(
+                stop_codes.eq(0) & committed.ge(max_new_tokens),
+                torch.full_like(stop_codes, 3),
+                stop_codes,
+            )
+            if generation_config.max_denoising_steps is not None:
+                stop_codes = torch.where(
+                    stop_codes.eq(0)
+                    & denoise_steps.ge(generation_config.max_denoising_steps),
+                    torch.full_like(stop_codes, 4),
+                    stop_codes,
+                )
+            stop_codes = torch.where(
+                stop_codes.eq(0) & denoise_steps.ge(max_iterations),
+                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())
+        sequences = torch.cat((input_ids, generated[:, :output_width]), dim=-1)
+        if streamer is not None:
+            streamer.end()
+        reason_names = {
+            1: "turn_end",
+            2: "eos",
+            3: "max_new_tokens",
+            4: "max_denoising_steps",
+            5: "episode_watchdog",
+        }
+        stop_reasons = tuple(
+            reason_names.get(code, "unknown") for code in stop_codes.detach().cpu().tolist()
+        )
+        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).flatten(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:
+                return value
+            item = value[0].item()
+            return float(item) if floating else int(item)
+
+        return ModilifyMk2GenerationOutput(
+            sequences=sequences,
+            generated_lengths=committed.clone(),
+            tokens_per_forward=tokens_per_forward,
+            past_key_values=past_key_values,
+            stop_reason=stop_reasons[0] if batch_size == 1 else stop_reasons,
+            committed_tokens=scalar_or_tensor(committed),
+            denoise_steps=scalar_or_tensor(denoise_steps),
+            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",
+]
+
+
+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)

inference.py

@@ -0,0 +1,43 @@
+"""Run the published Transformers model locally or from Hugging Face."""
+from __future__ import annotations
+import argparse
+from pathlib import Path
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", default=str(Path(__file__).resolve().parent))
+    parser.add_argument("--prompt", default="你好，请介绍一下自己。")
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda", "mps"), default="auto")
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--max-denoising-steps", type=int, default=256)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--no-thinking", action="store_true")
+    parser.add_argument("--stream", action="store_true")
+    args = parser.parse_args()
+    device = args.device
+    if device == "auto":
+        device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model, trust_remote_code=True, dtype="auto", device_map=device,
+        attn_implementation="sdpa",
+    ).eval()
+    inputs = tokenizer.apply_chat_template(
+        [{"role": "user", "content": args.prompt}],
+        enable_thinking=not args.no_thinking, add_generation_prompt=True,
+        return_tensors="pt", return_dict=True,
+    ).to(model.device)
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) if args.stream else None
+    with torch.inference_mode():
+        output = model.generate(**inputs, max_new_tokens=args.max_new_tokens,
+                                max_denoising_steps=args.max_denoising_steps,
+                                seed=args.seed, streamer=streamer)
+    if streamer is None:
+        print(tokenizer.decode(output.sequences[0, inputs.input_ids.shape[1]:], skip_special_tokens=True))
+
+
+if __name__ == "__main__":
+    main()

latent_deliberation.py

@@ -0,0 +1,827 @@
+"""Schema25 Torch GDN2 trajectory state, processor, and memory readers."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+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
+
+
+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(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    *,
+    attn_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run latent-memory attention reductions in FP32, then restore dtype.
+
+    These attention maps are small compared with the frozen decoder, while
+    their outputs feed recurrent trajectory and persistent-memory paths.  A
+    BF16 reduction error therefore compounds across denoise/commit steps and
+    is much more expensive than the modest FP32 workspace.
+    """
+
+    output_dtype = query.dtype
+    stable_mask = attn_mask
+    if stable_mask is not None and stable_mask.is_floating_point():
+        stable_mask = stable_mask.float()
+    query_fp32 = query.float()
+    key_fp32 = key.float()
+    use_batched_mm = query.ndim == 4 and key.ndim == 4 and value.ndim == 4
+    if use_batched_mm:
+        query_length = int(query.shape[-2])
+        key_length = int(key.shape[-2])
+        scores = torch.bmm(
+            query_fp32.reshape(-1, query_length, query.shape[-1]),
+            key_fp32.reshape(-1, key_length, key.shape[-1]).transpose(1, 2),
+        ).reshape(*query.shape[:-2], query_length, key_length)
+    else:
+        scores = torch.matmul(query_fp32, key_fp32.transpose(-2, -1))
+    scores = scores / math.sqrt(max(query.shape[-1], 1))
+    if stable_mask is not None:
+        if stable_mask.dtype == torch.bool:
+            scores = scores.masked_fill(~stable_mask, _sdpa_mask_value(torch.float32))
+        else:
+            scores = scores + stable_mask
+    probabilities = torch.softmax(scores, dim=-1)
+    # All-masked rows are uniform under a finite mask, but keep a NaN
+    # barrier for any remaining -inf path that MPS softmax cannot invert.
+    probabilities = torch.nan_to_num(probabilities, nan=0.0)
+    if use_batched_mm:
+        output = torch.bmm(
+            probabilities.reshape(-1, query.shape[-2], key.shape[-2]),
+            value.float().reshape(-1, key.shape[-2], value.shape[-1]),
+        ).reshape(*query.shape[:-2], query.shape[-2], value.shape[-1])
+    else:
+        output = torch.matmul(probabilities, value.float())
+    return output.to(dtype=output_dtype)
+
+
+@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 | None = None
+
+    @classmethod
+    def empty(
+        cls,
+        *,
+        batch_size: int,
+        canvas_length: int,
+        latent_dim: int,
+        memory_slots: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        enable_gdn2: bool = False,
+    ) -> "LatentDeliberationState":
+        persistent = (
+            torch.zeros(batch_size, 16, 128, 128, device=device, dtype=torch.float32)
+            if enable_gdn2 else None
+        )
+        return cls(
+            memory_slots=(persistent if persistent is not None else torch.zeros(
+                batch_size, memory_slots, latent_dim, device=device, dtype=dtype
+            )),
+            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),
+            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),
+            ) if enable_gdn2 else None,
+        )
+
+    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(),
+            gdn2=None if self.gdn2 is None else self.gdn2.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=None if self.gdn2 is None else self.gdn2.shift(
+                torch.full((self.confidence.shape[0],), committed,
+                           device=self.confidence.device, dtype=torch.long)
+            ),
+        )
+
+    def advance_ring(
+        self,
+        commit_lengths: torch.Tensor | int,
+        head: torch.Tensor | int,
+        *,
+        entropy_fill_value: float = 0.0,
+        reset_clocks_mask: torch.BoolTensor | None = None,
+    ) -> "LatentDeliberationState":
+        """Advance ring buffer state without physical tensor shifts, invalidating overwritten slots."""
+        batch, canvas_length = self.confidence.shape
+        device = self.confidence.device
+        if isinstance(commit_lengths, int):
+            lengths = torch.full((batch, 1), commit_lengths, device=device, dtype=torch.long)
+        else:
+            lengths = commit_lengths.to(device=device, dtype=torch.long).view(batch, 1)
+        if isinstance(head, int):
+            head_tensor = torch.full((batch, 1), head, device=device, dtype=torch.long)
+        else:
+            head_tensor = head.to(device=device, dtype=torch.long).view(batch, 1)
+        positions = torch.arange(canvas_length, device=device)[None, :]
+        old_offsets = (positions - head_tensor) % canvas_length
+        overwritten = old_offsets < lengths
+        fill_entropy = torch.as_tensor(entropy_fill_value, device=device, dtype=self.entropy.dtype)
+        committed = lengths.squeeze(1).gt(0) if reset_clocks_mask is None else reset_clocks_mask
+        return LatentDeliberationState(
+            memory_slots=self.memory_slots.clone(),
+            confidence=torch.where(overwritten, torch.zeros_like(self.confidence), self.confidence),
+            entropy=torch.where(overwritten, fill_entropy, self.entropy),
+            age=torch.where(overwritten, torch.zeros_like(self.age), self.age),
+            token_changed=torch.where(overwritten, torch.zeros_like(self.token_changed), self.token_changed),
+            confidence_delta=torch.where(overwritten, torch.zeros_like(self.confidence_delta), self.confidence_delta),
+            entropy_delta=torch.where(overwritten, torch.zeros_like(self.entropy_delta), self.entropy_delta),
+            ponder_steps=torch.where(committed, 0, self.ponder_steps).to(torch.int32),
+            stagnation_steps=torch.where(committed, 0, self.stagnation_steps).to(torch.int32),
+            gdn2=None if self.gdn2 is None else self.gdn2.clear_refill(
+                lengths.squeeze(1), head_tensor.squeeze(1)
+            ),
+        )
+
+
+@dataclass
+class LatentProcessorOutput:
+    context: torch.Tensor
+    working_state: torch.Tensor
+    state: LatentDeliberationState
+
+
+def advance_trajectory_clocks(
+    ponder_steps: torch.Tensor,
+    stagnation_steps: torch.Tensor,
+    *,
+    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
+    ):
+        raise ValueError("Trajectory clock inputs must share shape [batch].")
+    committed = commit_lengths.gt(0)
+    waiting = active_rows & ~committed
+    next_ponder = torch.where(
+        committed, torch.zeros_like(ponder_steps), ponder_steps + waiting.to(torch.int32)
+    )
+    next_stagnation = torch.where(
+        committed,
+        torch.zeros_like(stagnation_steps),
+        stagnation_steps + waiting.to(torch.int32),
+    )
+    return next_ponder.to(torch.int32), next_stagnation.to(torch.int32)
+
+
+def should_force_trajectory_jump(
+    stagnation_steps: torch.Tensor,
+    *,
+    progress_scores: torch.Tensor | None = None,
+    min_progress: float = 0.0,
+    stagnation_threshold: int,
+    ponder_steps: torch.Tensor | None = None,
+    max_ponder_steps: int | None = None,
+) -> torch.BoolTensor:
+    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:
+        jump = jump | ponder_steps.ge(max_ponder_steps)
+    return jump.to(torch.bool)
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1.0e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    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 * self.weight.float()).to(dtype=hidden.dtype)
+
+
+class _SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden: int) -> None:
+        super().__init__()
+        self.gate = nn.Linear(dim, hidden, bias=False)
+        self.up = nn.Linear(dim, hidden, bias=False)
+        self.down = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(hidden)) * self.up(hidden))
+
+
+class _RankAttention(nn.Module):
+    """Sequence attention in a rank-``kv_rank`` subspace, then map back to ``dim``."""
+
+    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.k_proj = nn.Linear(dim, kv_rank, bias=False)
+        self.v_proj = nn.Linear(dim, kv_rank, bias=False)
+        self.o_proj = nn.Linear(kv_rank, dim, bias=False)
+        self.q_norm = _RMSNorm(dim)
+        self.k_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
+        key_len = keys.shape[1]
+        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)
+        keys = self.k_proj(self.k_norm(keys)).view(batch, key_len, heads, head_dim).transpose(1, 2)
+        values = self.v_proj(values).view(batch, key_len, heads, head_dim).transpose(1, 2)
+        mask = attn_mask
+        if mask is not None and mask.ndim == 2:
+            mask = mask.view(1, 1, queries, key_len)
+        elif mask is not None and mask.ndim == 3:
+            mask = mask.unsqueeze(1)
+        context = _fp32_scaled_dot_product_attention(
+            query, keys, values, attn_mask=mask
+        )
+        context = context.transpose(1, 2).reshape(batch, queries, self.kv_rank)
+        return self.o_proj(context)
+
+
+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.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_readers: int,
+        memory_dim: int,
+        kv_rank: int,
+        *,
+        relative_bias: bool = False,
+        address_with_identity: bool = False,
+        max_relative_span: int = 256,
+    ) -> None:
+        super().__init__()
+        if num_readers > 0:
+            if kv_rank % num_heads:
+                raise ValueError("Memory bus rank must be divisible by heads.")
+            if hidden_size % num_heads:
+                raise ValueError("Memory bus hidden size must be divisible by heads.")
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_readers = num_readers
+        self.kv_rank = kv_rank
+        self.head_dim = kv_rank // num_heads if num_heads else kv_rank
+        self.relative_bias = relative_bias
+        self.address_with_identity = address_with_identity
+        self.memory_norm = _RMSNorm(memory_dim)
+        self.address_norm = _RMSNorm(memory_dim)
+        self.memory_to_hidden = (
+            nn.Identity()
+            if memory_dim == hidden_size
+            else nn.Linear(memory_dim, hidden_size, bias=False)
+        )
+        self.k_proj = nn.Linear(hidden_size, kv_rank, bias=False)
+        self.v_proj = nn.Linear(hidden_size, kv_rank, bias=False)
+        self.q_norm = _RMSNorm(hidden_size)
+        self.q_proj = nn.ModuleList(
+            [nn.Linear(hidden_size, kv_rank, bias=False) for _ in range(num_readers)]
+        )
+        self.o_proj = nn.ModuleList(
+            [nn.Linear(kv_rank, hidden_size, bias=False) for _ in range(num_readers)]
+        )
+        self.alpha = nn.Parameter(torch.zeros(max(num_readers, 1), max(num_heads, 1)))
+        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,
+        memory: torch.Tensor,
+        slot_identity: torch.Tensor | None = None,
+    ) -> 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.")
+            mapped_keys = self.memory_to_hidden(self.address_norm(memory + slot_identity))
+            mapped_values = self.memory_to_hidden(self.memory_norm(memory))
+        else:
+            mapped_keys = mapped_values = self.memory_to_hidden(self.memory_norm(memory))
+        batch, slots, _dim = mapped_keys.shape
+        heads = self.num_heads
+        head_dim = self.head_dim
+        keys = self.k_proj(mapped_keys).view(batch, slots, heads, head_dim).transpose(1, 2)
+        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,
+        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)
+        return self.rel_bias[:, rel].to(dtype=dtype)
+
+    def read(
+        self,
+        hidden: torch.Tensor,
+        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, 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)
+        return hidden + self.o_proj[reader_index](context)
+
+    @torch.no_grad()
+    def reset_identity_parameters(self) -> None:
+        self.alpha.zero_()
+        self.rel_bias.zero_()
+
+
+def slice_latent_state(
+    state: LatentDeliberationState, rows: slice | torch.Tensor
+) -> LatentDeliberationState:
+    return LatentDeliberationState(
+        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=None if state.gdn2 is None else GDN2TrajectoryState(
+            cells=state.gdn2.cells[rows], row=state.gdn2.row[rows],
+            persistent=state.gdn2.persistent[rows], seen=state.gdn2.seen[rows],
+        ),
+    )
+
+
+def cat_latent_states(states: Sequence[LatentDeliberationState]) -> LatentDeliberationState:
+    def cat(name: str) -> torch.Tensor:
+        return torch.cat([getattr(state, name) for state in states], dim=0)
+
+    return LatentDeliberationState(
+        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=None if states[0].gdn2 is None else 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),
+        ),
+    )
+
+
+def infer_commit_reason(
+    commit_lengths: torch.Tensor,
+    *,
+    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."""
+
+    reasons = torch.full(
+        commit_lengths.shape,
+        COMMIT_REASON_NONE,
+        device=commit_lengths.device,
+        dtype=torch.long,
+    )
+    committed = commit_lengths.gt(0)
+    default = (
+        COMMIT_REASON_TRAINING_RANDOM if training_random else COMMIT_REASON_NORMAL
+    )
+    reasons = torch.where(committed, torch.full_like(reasons, default), reasons)
+    if jump_rows is not None:
+        reasons = torch.where(
+            committed & jump_rows.to(dtype=torch.bool),
+            torch.full_like(reasons, COMMIT_REASON_FORCED_JUMP),
+            reasons,
+        )
+    if commit_token_ids is not None and terminal_token_ids:
+        positions = torch.arange(
+            commit_token_ids.shape[1], device=commit_token_ids.device
+        )[None, :]
+        selected = positions.lt(commit_lengths[:, None])
+        terminal = torch.zeros_like(committed)
+        for token_id in terminal_token_ids:
+            terminal |= (commit_token_ids.eq(int(token_id)) & selected).any(dim=-1)
+        reasons = torch.where(
+            committed & terminal,
+            torch.full_like(reasons, COMMIT_REASON_TERMINAL),
+            reasons,
+        )
+    return reasons
+
+
+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))
+        self.enabled = False
+        self.freeze()
+
+    def freeze(self) -> None:
+        self.enabled = False
+        self.alpha.requires_grad_(False)
+
+    def unfreeze(self) -> None:
+        self.enabled = self.num_readers > 0
+        self.alpha.requires_grad_(self.enabled)
+
+    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 or (self.training and not self.enabled):
+            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)
+        self.activation_checkpointing = True
+
+    def reset_identity_parameters(self) -> None:
+        self.working_memory_bus.reset_identity_parameters()
+        self.persistent_memory_bus.reset_identity_parameters()
+
+
+    def logical_seen(self, state: LatentDeliberationState,
+                     canvas_head: torch.Tensor | None) -> torch.Tensor:
+        seen = state.gdn2.seen
+        if canvas_head is None:
+            return seen
+        batch, canvas = seen.shape
+        index = (canvas_head[:, None] + torch.arange(canvas, device=seen.device)[None, :]) % canvas
+        return seen.gather(1, index)
+
+    def forward(self, *, token_embeddings: torch.Tensor, confidence: torch.Tensor,
+                entropy: torch.Tensor, state: LatentDeliberationState,
+
+                canvas_head: torch.Tensor | None = None) -> LatentProcessorOutput:
+        if state.gdn2 is None:
+            raise RuntimeError("Schema25 GDN2 state is missing.")
+        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, hidden, next_state)
+
+    def observe_state(self, state: LatentDeliberationState,
+                      heavy: torch.Tensor, working: torch.Tensor,
+                      live: torch.Tensor, head: torch.Tensor) -> LatentDeliberationState:
+        if state.gdn2 is None:
+            raise RuntimeError("Schema25 GDN2 state is missing.")
+        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,
+                     prefix_lengths: torch.Tensor | None = None,
+                     commit_reason: torch.Tensor | None = None,
+                     canvas_head: torch.Tensor | None = None):
+        del prefix_lengths
+        batch, canvas, width = working_state.shape
+        count = int(commit_lengths.max().item())
+        if count <= 0:
+            zero = memory.new_zeros(())
+            return memory, {"gate_mean": zero, "gate_max": zero,
+                            "gate_gt_01": zero, "gate_gt_05": zero,
+                            "delta_norm_mean": zero}
+        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)
+        delta = written - memory
+        zero = memory.new_zeros(())
+        return written, {"gate_mean": zero, "gate_max": zero,
+                         "gate_gt_01": zero, "gate_gt_05": zero,
+                         "delta_norm_mean": delta.float().norm(dim=(-2, -1)).mean()}
+
+
+__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",
+    "advance_trajectory_clocks",
+    "cat_latent_states",
+    "infer_commit_reason",
+    "should_force_trajectory_jump",
+    "slice_latent_state",
+]

model-00001-of-00016.safetensors

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

model-00002-of-00016.safetensors

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

model-00003-of-00016.safetensors

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

model-00004-of-00016.safetensors

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

model-00005-of-00016.safetensors

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

model-00006-of-00016.safetensors

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

model-00007-of-00016.safetensors

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

model-00008-of-00016.safetensors

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

model-00009-of-00016.safetensors

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

model-00010-of-00016.safetensors

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

model-00011-of-00016.safetensors

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

model-00012-of-00016.safetensors

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

model-00013-of-00016.safetensors

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

model-00014-of-00016.safetensors

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

model-00015-of-00016.safetensors

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

model-00016-of-00016.safetensors

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

modeling_modilify_mk2.py

@@ -0,0 +1,955 @@
+"""Text-only ModilifyMk2 model with recurrent latent deliberation."""
+
+from __future__ import annotations
+
+import copy
+from collections.abc import Sequence
+from dataclasses import dataclass, replace
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import BaseModelOutputWithPast
+
+from transformers.masking_utils import (
+    ALL_MASK_ATTENTION_FUNCTIONS,
+    bidirectional_mask_function,
+)
+from transformers.models.diffusion_gemma import (
+    DiffusionGemmaDecoderModel,
+    DiffusionGemmaEncoderTextModel,
+    DiffusionGemmaPreTrainedModel,
+)
+from transformers.models.diffusion_gemma.modeling_diffusion_gemma import (
+    DiffusionGemmaRMSNorm,
+    DiffusionGemmaTextRouter,
+)
+
+from .mps_ops import register_mps_backends
+from .adapters import LoRAConfig, inject_lora
+from .configuration_modilify_mk2 import ModilifyMk2Config, DENOISE_TEMPERATURE
+from .generation_modilify_mk2 import ModilifyMk2GenerationConfig, ModilifyMk2GenerationMixin
+from .latent_deliberation import LatentDeliberationState, LatentProcessorOutput, _sdpa_mask_value
+from .latent_deliberation import LatentDeliberationTransformer
+from .vocab_ops import (
+    chunked_vocab_statistics,
+    chunked_vocab_training_statistics,
+)
+
+
+register_mps_backends()
+
+
+@dataclass
+class ModilifyMk2DecoderOutput(BaseModelOutputWithPast):
+    token_embeddings: torch.FloatTensor | None = None
+    latent_residual_diagnostics: dict[str, torch.Tensor] | None = None
+
+
+@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:
+    logits: torch.FloatTensor | None
+    heavy_hidden_state: torch.FloatTensor
+    next_latent_state: LatentDeliberationState
+    gold_nll: torch.FloatTensor | None = None
+    loss: torch.FloatTensor | None = None
+    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
+    heavy_context_leaf: torch.FloatTensor | None = None
+    temporal_memory_slots: torch.FloatTensor | None = None
+    heavy_memory_leaf: torch.FloatTensor | None = None
+    working_state: torch.FloatTensor | None = None
+    latent_residual_diagnostics: dict[str, torch.Tensor] | None = None
+    temperature_gold_nll: torch.FloatTensor | None = None
+    proposal: torch.LongTensor | None = None
+    proposal_confidence: torch.FloatTensor | None = None
+    sampled_nll: 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.loss, 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)
+
+    def __getitem__(self, key: str | int) -> Any:
+        return getattr(self, key) if isinstance(key, str) else self.to_tuple()[key]
+
+
+class ModilifyMk2TextRouter(DiffusionGemmaTextRouter):
+    """Official router parameters with fp32 softmax so top-k weights stay finite."""
+
+    def __init__(self, config: Any) -> None:
+        super().__init__(config)
+        self.norm = DiffusionGemmaRMSNorm(self.hidden_size, eps=self.eps, with_scale=False)
+
+    def forward(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden_states = self.norm(hidden_states)
+        hidden_states = hidden_states * self.scale * self.scalar_root_size
+        expert_scores = self.proj(hidden_states)
+        # Match the official DiffusionGemma router: bf16 softmax underflows to
+        # all-zero top-k rows, then 0/0 NaNs the MoE residual and every loss.
+        router_probabilities = F.softmax(expert_scores, dim=-1, dtype=torch.float32)
+        top_k_weights, top_k_index = torch.topk(
+            router_probabilities,
+            k=self.config.top_k_experts,
+            dim=-1,
+        )
+        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True).clamp_min(
+            torch.finfo(torch.float32).tiny
+        )
+        top_k_weights = top_k_weights * self.per_expert_scale[top_k_index]
+        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 SDPA masks boolean and unmask fully-blocked query rows.
+
+    Transformers' SDPA mask is True=attend.  A fully-False query row makes
+    softmax(-inf) in the backward.  Converted additive 8k masks also double
+    decoder workspace on long prefixes, so stay in bool and only apply the
+    torch<2.5 empty-row workaround (pytorch/pytorch#110213).
+    """
+
+    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 DiffusionGemmaTextRouter:
+            replacement = ModilifyMk2TextRouter(child.config)
+            replacement.load_state_dict(child.state_dict())
+            setattr(module, name, replacement)
+        else:
+            install_modilify_mk2_trunk_semantics(child)
+
+
+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."""
+
+    config_class = ModilifyMk2Config
+    latent_residual_rms_ratio_cap = 0.5
+
+    def __init__(self, config: ModilifyMk2Config):
+        super().__init__(config)
+        self.gradient_checkpointing = False
+        install_modilify_mk2_trunk_semantics(self)
+
+    @staticmethod
+    def create_diffusion_decoder_attention_mask(
+        config: Any,
+        inputs_embeds: torch.Tensor,
+        past_key_values: Cache,
+        decoder_attention_mask: torch.Tensor | dict | None = None,
+    ) -> dict[str, torch.Tensor | None]:
+        """Official mask builder without the all-True sliding-window skip."""
+
+        if past_key_values is None:
+            raise ValueError(
+                "The diffusion mask requires `past_key_values` to construct the next attention mask correctly"
+            )
+        if (
+            decoder_attention_mask is None
+            or config._attn_implementation
+            not in ALL_MASK_ATTENTION_FUNCTIONS._global_mapping
+        ):
+            return {"full_attention": None, "sliding_attention": None}
+        if isinstance(decoder_attention_mask, dict) and all(
+            mask.ndim == 4 for mask in decoder_attention_mask.values()
+        ):
+            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]
+        q_offset = past_key_values.get_seq_length()
+        if isinstance(q_offset, torch.Tensor):
+            q_offset = q_offset.to(inputs_embeds.device)
+        additional_kv_length = (
+            getattr(config, "canvas_length", 0) if past_key_values.is_compileable else 0
+        )
+        mask_mapping: dict[str, torch.Tensor | None] = {}
+        for layer_pattern in set(text_config.layer_types):
+            layer_idx = past_key_values.is_sliding.index(
+                layer_pattern == "sliding_attention"
+            )
+            kv_length, kv_offset = past_key_values.get_mask_sizes(q_length, layer_idx)
+            kv_length += additional_kv_length
+            if layer_pattern == "sliding_attention" and past_key_values.is_compileable:
+                sliding_layer = past_key_values.layers[layer_idx]
+                max_length = sliding_layer.get_max_length() + additional_kv_length
+                if kv_length >= max_length:
+                    kv_length = max_length
+            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,
+            )
+        return mask_mapping
+
+    def merge_latent_context(
+        self,
+        token_embeddings: torch.Tensor,
+        latent_context: torch.Tensor | None,
+        *,
+        collect_diagnostics: bool = True,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """Map latent context through the frozen native self-conditioning bridge."""
+
+        if latent_context is None:
+            context = torch.zeros_like(token_embeddings)
+        else:
+            if latent_context.shape != token_embeddings.shape:
+                raise ValueError("Temporal context must match canvas hidden-state shape.")
+            context = latent_context.to(token_embeddings)
+        mapper = self.self_conditioning
+        normalized_context = mapper.pre_norm(context)
+        mapped_context = mapper.down_proj(
+            mapper.act_fn(mapper.gate_proj(normalized_context))
+            * 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()
+        )
+        residual_cap = self.latent_residual_rms_ratio_cap * token_rms_per_token
+        soft_cap_scale = residual_cap / torch.sqrt(
+            mapped_token_energy + residual_cap.square() + 1.0e-12
+        )
+        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
+
+    def _run_stack(
+        self,
+        inputs_embeds: torch.Tensor,
+        *,
+        past_key_values: Cache | None,
+        decoder_attention_mask: torch.Tensor | dict | None,
+        decoder_position_ids: torch.LongTensor | None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        if decoder_position_ids is None:
+            prefix = past_key_values.get_seq_length(layer_idx=0) if past_key_values is not None else 0
+            decoder_position_ids = torch.arange(
+                prefix, prefix + inputs_embeds.shape[1], device=inputs_embeds.device
+            ).unsqueeze(0)
+        if not isinstance(mask_mapping := decoder_attention_mask, dict):
+            mask_mapping = self.create_diffusion_decoder_attention_mask(
+                config=self.text_config,
+                inputs_embeds=inputs_embeds,
+                past_key_values=past_key_values,
+                decoder_attention_mask=decoder_attention_mask,
+            )
+        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) or kwargs.pop("memory_bus", None)
+        memory_slots = kwargs.pop("memory_slots", None)
+        slot_identity = kwargs.pop("slot_identity", None)
+        cache = past_key_values
+        if self.gradient_checkpointing and self.training and cache is not None:
+            cache = copy.copy(cache)
+            cache.layers = [copy.copy(layer) for layer in cache.layers]
+        hidden = inputs_embeds
+        positions = {
+            layer_type: self.rotary_emb(hidden, decoder_position_ids, layer_type)
+            for layer_type in self.unique_layer_types
+        }
+        working_kv = None
+        if working_bus is not None and working_state is not None:
+            working_kv = working_bus.prepare_kv(working_state)
+        memory_kv = None
+        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):
+            layer = self.layers[index]
+            layer_type = self.text_config.layer_types[index]
+            if self.gradient_checkpointing and self.training:
+                def layer_forward(
+                    values: torch.Tensor,
+                    selected: nn.Module = layer,
+                    selected_type: str = layer_type,
+                ) -> torch.Tensor:
+                    return selected(
+                        values,
+                        position_embeddings=positions[selected_type],
+                        attention_mask=mask_mapping[selected_type],
+                        position_ids=decoder_position_ids,
+                        past_key_values=cache,
+                        **kwargs,
+                    )
+
+                hidden = self._gradient_checkpointing_func(layer_forward, hidden)
+            else:
+                hidden = layer(
+                    hidden,
+                    position_embeddings=positions[layer_type],
+                    attention_mask=mask_mapping[layer_type],
+                    position_ids=decoder_position_ids,
+                    past_key_values=cache,
+                    **kwargs,
+                )
+            if layer_type == "full_attention":
+                if (
+                    working_kv is not None
+                    and working_reader < working_bus.num_readers
+                ):
+                    hidden = working_bus.read(
+                        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 < persistent_bus.num_readers
+                ):
+                    hidden = persistent_bus.read(
+                        hidden, reader_index, memory_kv[0], memory_kv[1]
+                    )
+                    reader_index += 1
+        return self.norm(hidden)
+
+    def forward(
+        self,
+        decoder_input_ids: torch.LongTensor,
+        past_key_values: Cache | None = None,
+        decoder_token_embeddings: torch.FloatTensor | None = None,
+        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_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:
+        if "use_cache" in kwargs:
+            raise ValueError("The diffusion decoder always reads the supplied cache.")
+        if decoder_token_embeddings is None:
+            token_embeddings = self.embed_tokens(decoder_input_ids)
+        else:
+            token_embeddings = decoder_token_embeddings
+            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(
+            token_embeddings,
+            context_embeddings if temporal_context_embeddings is not None else None,
+            collect_diagnostics=collect_latent_diagnostics,
+        )
+        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_slots=memory_slots,
+            working_bus=working_bus,
+            working_state=working_state,
+            persistent_bus=persistent_bus,
+            slot_identity=slot_identity,
+            **kwargs,
+        )
+        return ModilifyMk2DecoderOutput(
+            last_hidden_state=hidden,
+            past_key_values=past_key_values,
+            token_embeddings=token_embeddings,
+            latent_residual_diagnostics=diagnostics,
+        )
+
+
+class ModilifyMk2Model(DiffusionGemmaPreTrainedModel):
+    config_class = ModilifyMk2Config
+    _tied_weights_keys = {
+        "encoder.language_model.norm.weight": "decoder.norm.weight",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*weight": r"decoder.layers\.(?:[^.]+\.)*weight",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*scale": r"decoder.layers\.(?:[^.]+\.)*scale",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*per_expert_scale": r"decoder.layers\.(?:[^.]+\.)*per_expert_scale",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*gate_up_proj": r"decoder.layers\.(?:[^.]+\.)*gate_up_proj",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*down_proj": r"decoder.layers\.(?:[^.]+\.)*down_proj",
+        "encoder.language_model.embed_tokens.weight": "decoder.embed_tokens.weight",
+    }
+
+    def __init__(self, config: ModilifyMk2Config):
+        super().__init__(config)
+        self.encoder = ModilifyMk2EncoderModel(config)
+        self.decoder = ModilifyMk2DecoderModel(config)
+        install_modilify_mk2_trunk_semantics(self)
+        self.post_init()
+
+    def get_encoder(self):
+        return self.encoder
+
+    def get_decoder(self):
+        return self.decoder
+
+    def get_input_embeddings(self):
+        return self.encoder.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.encoder.set_input_embeddings(value)
+        self.decoder.embed_tokens = value
+
+    def forward(
+        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,
+        decoder_token_embeddings: torch.FloatTensor | None = None,
+        temporal_context_embeddings: torch.FloatTensor | None = None,
+        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_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,
+    ) -> ModilifyMk2ModelOutput:
+        encoder_hidden = None
+        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,
+                **kwargs,
+            )
+            past_key_values = encoded.past_key_values
+            if return_encoder_outputs:
+                encoder_hidden = encoded.last_hidden_state
+        elif past_key_values is None:
+            raise ValueError("Either `input_ids` or `past_key_values` is required.")
+        decoded = self.decoder(
+            decoder_input_ids=decoder_input_ids,
+            decoder_token_embeddings=decoder_token_embeddings,
+            past_key_values=past_key_values,
+            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_slots=memory_slots,
+            working_bus=working_bus,
+            working_state=working_state,
+            persistent_bus=persistent_bus,
+            slot_identity=slot_identity,
+            **kwargs,
+        )
+        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,
+        )
+
+
+class ModilifyMk2ForBlockDiffusion(DiffusionGemmaPreTrainedModel, ModilifyMk2GenerationMixin):
+    config_class = ModilifyMk2Config
+    _tied_weights_keys = {"lm_head.weight": "model.decoder.embed_tokens.weight"}
+    generation_config_class = ModilifyMk2GenerationConfig
+    _no_split_modules = ["DiffusionGemmaDecoderLayer", "LatentDeliberationTransformer"]
+    supports_gradient_checkpointing = True
+
+    @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)
+        self.model = ModilifyMk2Model(config)
+        layer_types = tuple(getattr(config.text_config, "layer_types", None) or ())
+        self.latent_deliberation = LatentDeliberationTransformer(
+            hidden_size=config.text_config.hidden_size,
+            latent_dim=config.latent_dim,
+            ffn_dim=config.latent_ffn_dim,
+            num_layers=config.latent_num_layers,
+            num_heads=config.latent_num_heads,
+            local_attention_window=config.latent_local_attention_window,
+            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),
+            num_working_readers=(
+                sum(layer_type == "full_attention" for layer_type in layer_types)
+                if config.working_memory_bus else 0
+            ),
+            num_persistent_readers=(
+                sum(layer_type == "full_attention" for layer_type in layer_types)
+                if config.persistent_memory_bus else 0
+            ),
+            working_last_block_global=config.latent_working_last_block_global,
+            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
+        self.post_init()
+        install_modilify_mk2_trunk_semantics(self)
+        if getattr(config, "lora_config", None):
+            payload = dict(config.lora_config)
+            payload["target_modules"] = tuple(payload["target_modules"])
+            inject_lora(self, LoRAConfig(**payload))
+            # The same adapters serve the encoder and decoder, just like the
+            # tied backbone. Tell from_pretrained/save_pretrained about them.
+            self._tied_weights_keys = dict(type(self)._tied_weights_keys)
+            self._tied_weights_keys.update({
+                name: name.replace("model.encoder.language_model.", "model.decoder.", 1)
+                for name, _ in self.named_parameters(remove_duplicate=False)
+                if name.startswith("model.encoder.language_model.") and ".lora_" in name
+            })
+            self.all_tied_weights_keys = self.get_expanded_tied_weights_keys(all_submodels=True)
+            self.tie_weights(recompute_mapping=False)
+            config.text_config._experts_implementation = "modilify_mk2_lora"
+        # These 26 small parameters were trained and saved in FP32. Preserve
+        # their dtype even when from_pretrained is called with BF16/FP16.
+        self._keep_in_fp32_modules_strict = {
+            name for name, _ in self.named_parameters()
+            if name.startswith("latent_deliberation.") and (
+                name.endswith((".dt_bias", ".a_log")) or
+                (name.endswith(".weight") and "norm" in name.rsplit(".", 2)[-2])
+            )
+        }
+
+    def _finalize_logits(self, hidden: torch.Tensor) -> torch.Tensor:
+        logits = self.lm_head(hidden)
+        return torch.tanh(logits / self.final_logit_softcapping) * self.final_logit_softcapping
+
+    def _compute_gold_nll(
+        self, logits: torch.Tensor, labels: torch.LongTensor | None
+    ) -> torch.Tensor | None:
+        if labels is None:
+            return None
+        labels = labels.to(logits.device)
+        return F.cross_entropy(
+            logits.flatten(0, 1), labels.flatten(),
+            ignore_index=-100, reduction="none",
+        ).view_as(labels)
+
+    def _prepare_latent_context(
+        self,
+        decoder_input_ids: torch.LongTensor,
+        *,
+
+
+        confidence: torch.Tensor | None,
+        entropy: torch.Tensor | None,
+        age: torch.Tensor | None,
+        latent_state: LatentDeliberationState | None,
+        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, enable_gdn2=True,
+            )
+        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)
+        use_outer_checkpoint = (
+            self.training
+            and torch.is_grad_enabled()
+            and self.latent_deliberation.activation_checkpointing
+        )
+        if use_outer_checkpoint:
+            def run_latent_processor(
+                embeddings: torch.Tensor,
+                memory_slots: torch.Tensor,
+                working_cells: torch.Tensor,
+                working_row: torch.Tensor,
+                state_confidence: torch.Tensor,
+                state_entropy: torch.Tensor,
+                state_age: torch.Tensor,
+                state_token_changed: torch.Tensor,
+                state_confidence_delta: torch.Tensor,
+                state_entropy_delta: torch.Tensor,
+                state_ponder_steps: torch.Tensor,
+                state_stagnation_steps: torch.Tensor,
+
+
+            ) -> torch.Tensor:
+                checkpoint_state = LatentDeliberationState(
+                    memory_slots=memory_slots,
+                    confidence=state_confidence,
+                    entropy=state_entropy,
+                    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,
+                    gdn2=replace(
+                        latent_state.gdn2, cells=working_cells,
+                        row=working_row, persistent=memory_slots,
+                    ),
+                )
+                result = self.latent_deliberation(
+                    token_embeddings=embeddings,
+                    confidence=confidence,
+                    entropy=entropy,
+                    state=checkpoint_state,
+
+
+                    canvas_head=canvas_head,
+                )
+                return result.context
+
+            latent_context = checkpoint(
+                run_latent_processor,
+                token_embeddings,
+                latent_state.memory_slots,
+                latent_state.gdn2.cells,
+                latent_state.gdn2.row,
+                latent_state.confidence,
+                latent_state.entropy,
+                latent_state.age,
+                latent_state.token_changed,
+                latent_state.confidence_delta,
+                latent_state.entropy_delta,
+                latent_state.ponder_steps,
+                latent_state.stagnation_steps,
+
+
+                use_reentrant=False,
+                preserve_rng_state=True,
+            )
+            next_state = LatentDeliberationState(
+                memory_slots=latent_state.memory_slots,
+                confidence=confidence.to(dtype=torch.float32),
+                entropy=entropy.to(dtype=torch.float32),
+                age=latent_state.age,
+                token_changed=latent_state.token_changed,
+                confidence_delta=latent_state.confidence_delta,
+                entropy_delta=latent_state.entropy_delta,
+                ponder_steps=latent_state.ponder_steps,
+                stagnation_steps=latent_state.stagnation_steps,
+                gdn2=latent_state.gdn2,
+            )
+        else:
+            processed: LatentProcessorOutput = self.latent_deliberation(
+                token_embeddings=token_embeddings,
+                confidence=confidence,
+                entropy=entropy,
+                state=latent_state,
+
+
+                canvas_head=canvas_head,
+            )
+            latent_context = processed.context
+            next_state = processed.state
+        return (
+            latent_context,
+            next_state,
+            token_embeddings,
+        )
+
+    def forward(
+        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,
+
+
+        canvas_head: torch.Tensor | None = None,
+
+        decoder_attention_mask: torch.Tensor | dict | None = None,
+        decoder_position_ids: torch.LongTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        return_encoder_outputs: bool = True,
+        detach_temporal_context_for_heavy: bool = False,
+        compact_vocab: bool = False,
+        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:
+        (
+            latent_context,
+            next_state,
+            decoder_token_embeddings,
+        ) = self._prepare_latent_context(
+            decoder_input_ids,
+
+
+            confidence=previous_confidence,
+            entropy=previous_entropy,
+            age=token_age,
+            latent_state=latent_state,
+            canvas_head=canvas_head,
+        )
+        heavy_context = latent_context
+        heavy_context_leaf = None
+        temporal_memory_slots = next_state.memory_slots
+        heavy_memory_slots = temporal_memory_slots
+        heavy_memory_leaf = None
+        working_bus = self.latent_deliberation.working_memory_bus
+        persistent_bus = self.latent_deliberation.persistent_memory_bus
+        working_readers_active = (
+            working_bus.num_readers > 0
+            and (not working_bus.training or working_bus.enabled)
+        )
+        memory_readers_active = (
+            persistent_bus.num_readers > 0
+            and (not persistent_bus.training or persistent_bus.enabled)
+        )
+        if detach_temporal_context_for_heavy and latent_context.requires_grad:
+            # Clone so the heavy decoder cannot alias the temporal graph storage.
+            # `detach()` alone shares storage; an in-place decoder write would
+            # corrupt the later latent backward and can produce non-finite grads.
+            heavy_context_leaf = latent_context.detach().clone().requires_grad_(True)
+            heavy_context = heavy_context_leaf
+        if (
+            detach_temporal_context_for_heavy
+            and memory_readers_active
+            and temporal_memory_slots.requires_grad
+        ):
+            # Memory readers are part of the heavy decoder too.  Keep their
+            # backward separate from the recurrent temporal graph for the same
+            # reason as the canvas context above; both gradients are replayed
+            # together at the TBPTT boundary.
+            heavy_memory_leaf = (
+                temporal_memory_slots.detach().clone().requires_grad_(True)
+            )
+            heavy_memory_slots = heavy_memory_leaf
+        # The production LoRA path freezes the tied embedding table. Reuse its
+        # lookup there; a trainable table needs two independent autograd paths
+        # because heavy and temporal backward execute at different times.
+        reusable_decoder_embeddings = (
+            None
+            if self.model.decoder.embed_tokens.weight.requires_grad
+            else decoder_token_embeddings
+        )
+        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=reusable_decoder_embeddings,
+            temporal_context_embeddings=heavy_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=(
+                heavy_context,
+                # The Torch decoder keeps physical ring order and supplies
+                # logical positions through RoPE; its mask must stay physical.
+                next_state.gdn2.seen,
+            ) if working_readers_active else None,
+            working_canvas_head=canvas_head,
+            persistent_bus=persistent_bus if memory_readers_active else None,
+            memory_slots=heavy_memory_slots if memory_readers_active else None,
+            slot_identity=(
+                next_state.gdn2.seen
+                if memory_readers_active else None
+            ),
+            **kwargs,
+        )
+        latent_residual_diagnostics = outputs.latent_residual_diagnostics
+        proposal = proposal_confidence = token_entropy = None
+        greedy_proposal = greedy_confidence = None
+        temperature_gold_nll = None
+        sampled_nll = None
+        if compact_vocab:
+            logits = None
+            if labels is None:
+                (
+                    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=DENOISE_TEMPERATURE,
+                        chunk_size=self.config.vocab_chunk_size,
+                        repetition_token_mask=repetition_token_mask,
+                        repetition_penalty=repetition_penalty,
+                        sampling_generators=sampling_generators,
+                        top_k=getattr(self.config, "commit_top_k", 40),
+                        min_p=getattr(self.config, "commit_min_p", 0.05),
+                    )
+                )
+                gold_nll = None
+            else:
+                (
+                    gold_nll,
+                    temperature_gold_nll,
+                    proposal,
+                    proposal_confidence,
+                    token_entropy,
+                    greedy_proposal,
+                    greedy_confidence,
+                    sampled_nll,
+                ) = chunked_vocab_training_statistics(
+                    outputs.last_hidden_state,
+                    self.lm_head.weight,
+                    labels,
+                    return_sampled_nll=True,
+                    softcap=self.final_logit_softcapping,
+                    temperature=DENOISE_TEMPERATURE,
+                    chunk_size=self.config.vocab_chunk_size,
+                    top_k=getattr(self.config, "commit_top_k", 40),
+                    min_p=getattr(self.config, "commit_min_p", 0.05),
+                )
+        else:
+            logits = self._finalize_logits(outputs.last_hidden_state)
+            gold_nll = self._compute_gold_nll(logits, labels)
+        loss = None
+        if gold_nll is not None and labels is not None:
+            valid = labels.to(gold_nll.device).ne(-100)
+            loss = gold_nll.masked_fill(~valid, 0.0).sum() / valid.sum().clamp_min(1)
+        return ModilifyMk2BlockDiffusionOutput(
+            logits=logits, loss=loss,
+            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,
+            gold_nll=gold_nll,
+            temporal_context=latent_context,
+            heavy_context_leaf=heavy_context_leaf,
+            temporal_memory_slots=temporal_memory_slots,
+            heavy_memory_leaf=heavy_memory_leaf,
+
+            working_state=latent_context,
+            latent_residual_diagnostics=latent_residual_diagnostics,
+            temperature_gold_nll=temperature_gold_nll,
+            proposal=proposal,
+            proposal_confidence=proposal_confidence,
+            sampled_nll=sampled_nll,
+            token_entropy=token_entropy,
+            greedy_proposal=greedy_proposal,
+            greedy_confidence=greedy_confidence,
+        )
+
+
+__all__ = [
+    "ModilifyMk2BlockDiffusionOutput", "ModilifyMk2Config", "ModilifyMk2DecoderModel",
+    "ModilifyMk2EncoderModel", "ModilifyMk2ForBlockDiffusion", "ModilifyMk2Model",
+    "ModilifyMk2TextRouter", "install_modilify_mk2_trunk_semantics",
+]

mps_ops.py

@@ -0,0 +1,921 @@
+"""MPS-specific kernels that preserve the model's mathematical operations."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS, _grouped_linear
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+
+def attach_expert_lora(
+    experts: nn.Module,
+    *,
+    rank: int,
+    alpha: int,
+    dropout: float,
+) -> 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.")
+    if not 0.0 <= dropout < 1.0:
+        raise ValueError("Expert LoRA dropout must be in [0, 1).")
+    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_dropout = nn.Dropout(dropout) if dropout else nn.Identity()
+    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,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    final_hidden_states = torch.zeros_like(hidden_states)
+    with torch.no_grad():
+        expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
+        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:
+        expert_idx = expert_idx[0]
+        if expert_idx == self.num_experts:
+            continue
+        top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+        current_state = hidden_states[token_idx]
+        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]),
+                self.lora_gate_up_b[expert_idx],
+            )
+            gate_up = gate_up + lora_scaling * update
+        gate, up = gate_up.chunk(2, dim=-1)
+        current_hidden_states = self.act_fn(gate) * up
+        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]),
+                self.lora_down_b[expert_idx],
+            )
+            expert_output = expert_output + lora_scaling * update
+        current_hidden_states = expert_output * top_k_weights[token_idx, top_k_pos, None]
+        final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+    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)
+
+
+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
+
+
+_INDEPENDENT_EXPERT_LORA_METAL_SOURCE = r"""
+#include <metal_stdlib>
+using namespace metal;
+
+constant uint kRank = 8;
+constant uint kThreads = 256;
+
+kernel void expert_lora_forward_a(
+    device const bfloat* inputs [[buffer(0)]],
+    device const bfloat* weight_a [[buffer(1)]],
+    device const long* expert_ids [[buffer(2)]],
+    device bfloat* low_rank [[buffer(3)]],
+    constant uint& row_count [[buffer(4)]],
+    constant uint& input_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 a_base = (expert * kRank + rank) * input_dim;
+    const uint x_base = row * input_dim;
+    for (uint column = lane; column < input_dim; column += kThreads) {
+      value += float(inputs[x_base + column]) * float(weight_a[a_base + column]);
+    }
+    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) {
+      low_rank[row * kRank + rank] = bfloat(partial[rank]);
+    }
+  }
+}
+
+kernel void expert_lora_forward_b(
+    device const bfloat* low_rank [[buffer(0)]],
+    device const bfloat* weight_b [[buffer(1)]],
+    device const long* expert_ids [[buffer(2)]],
+    device bfloat* output [[buffer(3)]],
+    constant uint& row_count [[buffer(4)]],
+    constant uint& output_dim [[buffer(5)]],
+    uint index [[thread_position_in_grid]]) {
+  const uint total = row_count * output_dim;
+  if (index >= total) return;
+  const uint row = index / output_dim;
+  const uint column = index - row * output_dim;
+  const uint expert = uint(expert_ids[row]);
+  const uint b_base = (expert * output_dim + column) * kRank;
+  float value = 0.0f;
+  for (uint rank = 0; rank < kRank; ++rank) {
+    value += float(low_rank[row * kRank + rank])
+        * float(weight_b[b_base + rank]);
+  }
+  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);
+}
+"""
+
+
+_independent_expert_lora_metal_library = None
+
+
+def _get_independent_expert_lora_metal_library():
+    global _independent_expert_lora_metal_library
+    if _independent_expert_lora_metal_library is None:
+        _independent_expert_lora_metal_library = torch.mps.compile_shader(
+            _INDEPENDENT_EXPERT_LORA_METAL_SOURCE
+        )
+    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 _independent_expert_lora(
+    inputs: torch.Tensor,
+    weight_a: torch.Tensor,
+    weight_b: torch.Tensor,
+    expert_ids: torch.LongTensor,
+    counts: torch.LongTensor,
+) -> torch.Tensor:
+    """Dispatch the compact independent-expert operator without sharing A/B."""
+
+    if inputs.shape[0] == 0:
+        return inputs.new_empty((0, weight_b.shape[1]))
+    if (
+        inputs.device.type == "mps"
+        and inputs.dtype == torch.bfloat16
+        and weight_a.dtype == torch.bfloat16
+        and weight_b.dtype == torch.bfloat16
+        and weight_a.shape[1] == 8
+        and hasattr(torch.mps, "compile_shader")
+    ):
+        return _MetalIndependentExpertLoRA.apply(
+            inputs,
+            weight_a,
+            weight_b,
+            expert_ids,
+            counts,
+        )
+    return _IndependentExpertLoRA.apply(
+        inputs,
+        weight_a,
+        weight_b,
+        expert_ids,
+        counts,
+        int(counts.max().detach().cpu()),
+    )
+
+
+def mps_segmented_experts_forward(
+    self: torch.nn.Module,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    """Run GPU-grouped base experts plus independent Metal LoRA A/B."""
+    if hidden_states.device.type != "mps":
+        return lora_aware_eager_experts_forward(
+            self,
+            hidden_states,
+            top_k_index,
+            top_k_weights,
+        )
+
+    num_top_k = top_k_index.size(-1)
+    num_tokens, hidden_dim = hidden_states.shape
+    sample_weights = top_k_weights.reshape(-1)
+    expert_ids = top_k_index.reshape(-1)
+    expert_ids_g, permutation = torch.sort(expert_ids)
+    selected_hidden_g = hidden_states[permutation // num_top_k]
+    sample_weights_g = sample_weights[permutation]
+
+    counts = torch.bincount(
+        expert_ids_g,
+        minlength=self.num_experts,
+    )[: self.num_experts]
+    offsets = torch.cumsum(counts, dim=0, dtype=torch.int32)
+    gate_up_weights = self.gate_up_proj if self.has_gate else self.up_proj
+    gate_up_bias = None
+    if self.has_bias:
+        packed_bias = self.gate_up_proj_bias if self.has_gate else self.up_proj_bias
+        gate_up_bias = packed_bias[expert_ids_g]
+    gate_up_g = _grouped_linear(
+        selected_hidden_g,
+        gate_up_weights,
+        offsets,
+        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 = gate_up_g + gate_up_update * self.lora_scaling
+    elif hasattr(self, "lora_gate_up_a"):
+        gate_up_update = _independent_expert_lora(
+            selected_hidden_g,
+            self.lora_gate_up_a,
+            self.lora_gate_up_b,
+            expert_ids_g,
+            counts,
+        )
+        gate_up_g = gate_up_g + gate_up_update * self.lora_scaling
+    activated_g = (
+        self._apply_gate(gate_up_g)
+        if self.has_gate
+        else self.act_fn(gate_up_g)
+    )
+    down_bias = self.down_proj_bias[expert_ids_g] if self.has_bias else None
+    projected_g = _grouped_linear(
+        activated_g,
+        self.down_proj,
+        offsets,
+        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 = projected_g + down_update * self.lora_scaling
+    elif hasattr(self, "lora_down_a"):
+        down_update = _independent_expert_lora(
+            activated_g,
+            self.lora_down_a,
+            self.lora_down_b,
+            expert_ids_g,
+            counts,
+        )
+        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)
+
+
+_SDPA_MASK_FLOOR = -1.0e9
+_LONG_KV_CHUNK = 256
+_LONG_KV_THRESHOLD = 2048
+_original_sdpa_attention_forward = ALL_ATTENTION_FUNCTIONS.get_interface("sdpa", None)
+
+
+class _MpsGqaChunkedAttention(torch.autograd.Function):
+    """Exact GQA attention that never repeats long KV onto every query head.
+
+    Transformers' SDPA path `repeat_kv`s 2 global heads to 16 for full
+    attention.  On an 8k prefix that materializes ~2GB of K and V per layer;
+    mlp-moe retains all 5 full layers, and the second denoise still holds the
+    first latent graph.  MPS then returns NaN LoRA grads with a finite loss.
+    Saving compact GQA tensors and reducing K in chunks keeps the math
+    identical while staying inside unified memory.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: torch.Tensor | None,
+        scale: float,
+        chunk_size: int,
+    ) -> torch.Tensor:
+        query_fp32 = query.float()
+        key_fp32 = key.float()
+        value_fp32 = value.float()
+        batch, query_heads, queries, head_dim = query_fp32.shape
+        key_heads = int(key_fp32.shape[1])
+        key_len = int(key_fp32.shape[-2])
+        repeats = query_heads // key_heads
+        packed = batch * query_heads
+        query_flat = query_fp32.reshape(packed, queries, head_dim)
+        negative_inf = query_fp32.new_full((), _SDPA_MASK_FLOOR)
+        running_max = query_fp32.new_full((packed, queries, 1), float("-inf"))
+        running_sum = torch.zeros_like(running_max)
+        acc = torch.zeros(
+            packed, queries, head_dim,
+            device=query_fp32.device, dtype=query_fp32.dtype,
+        )
+        chunk = max(int(chunk_size), 1)
+
+        def _pack_kv(tensor: torch.Tensor, start: int, stop: int) -> torch.Tensor:
+            width = stop - start
+            packed_kv = tensor[:, :, start:stop, :].unsqueeze(2).expand(
+                batch, key_heads, repeats, width, head_dim
+            )
+            return packed_kv.reshape(packed, width, head_dim)
+
+        for start in range(0, key_len, chunk):
+            stop = min(start + chunk, key_len)
+            key_flat = _pack_kv(key_fp32, start, stop)
+            value_flat = _pack_kv(value_fp32, start, stop)
+            scores = torch.bmm(query_flat, key_flat.transpose(1, 2)) * scale
+            if attn_mask is not None:
+                mask_chunk = attn_mask[..., start:stop]
+                if mask_chunk.dtype == torch.bool:
+                    scores = scores.masked_fill(
+                        ~mask_chunk.expand(batch, query_heads, queries, stop - start).reshape(
+                            packed, queries, stop - start
+                        ),
+                        negative_inf,
+                    )
+                else:
+                    scores = scores + mask_chunk.float().expand(
+                        batch, query_heads, queries, stop - start
+                    ).reshape(packed, queries, stop - start)
+            block_max = scores.amax(dim=-1, keepdim=True)
+            new_max = torch.maximum(running_max, block_max)
+            weights = torch.exp(scores - new_max)
+            rescale = torch.exp(running_max - new_max)
+            running_sum = running_sum * rescale + weights.sum(dim=-1, keepdim=True)
+            acc = acc * rescale + torch.bmm(weights, value_flat)
+            running_max = new_max
+        output = acc / running_sum.clamp_min(torch.finfo(running_sum.dtype).tiny)
+        output = output.view(batch, query_heads, queries, head_dim)
+        ctx.save_for_backward(
+            query, key, value,
+            attn_mask if attn_mask is not None else query.new_empty(0),
+        )
+        ctx.scale = float(scale)
+        ctx.chunk_size = chunk
+        ctx.has_mask = attn_mask is not None
+        return output.to(dtype=query.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        query, key, value, mask_or_empty = ctx.saved_tensors
+        attn_mask = mask_or_empty if ctx.has_mask else None
+        scale = ctx.scale
+        chunk = ctx.chunk_size
+        query_fp32 = query.float()
+        key_fp32 = key.float()
+        value_fp32 = value.float()
+        grad_out = grad_output.float()
+        batch, query_heads, queries, head_dim = query_fp32.shape
+        key_heads = int(key_fp32.shape[1])
+        key_len = int(key_fp32.shape[-2])
+        repeats = query_heads // key_heads
+        packed = batch * query_heads
+        query_flat = query_fp32.reshape(packed, queries, head_dim)
+        grad_out_flat = grad_out.reshape(packed, queries, head_dim)
+        negative_inf = query_fp32.new_full((), _SDPA_MASK_FLOOR)
+        running_max = query_fp32.new_full((packed, queries, 1), float("-inf"))
+        running_sum = torch.zeros_like(running_max)
+
+        def _pack_kv(tensor: torch.Tensor, start: int, stop: int) -> torch.Tensor:
+            width = stop - start
+            packed_kv = tensor[:, :, start:stop, :].unsqueeze(2).expand(
+                batch, key_heads, repeats, width, head_dim
+            )
+            return packed_kv.reshape(packed, width, head_dim)
+
+        def _apply_mask(scores: torch.Tensor, start: int, stop: int) -> torch.Tensor:
+            if attn_mask is None:
+                return scores
+            mask_chunk = attn_mask[..., start:stop]
+            width = stop - start
+            if mask_chunk.dtype == torch.bool:
+                return scores.masked_fill(
+                    ~mask_chunk.expand(batch, query_heads, queries, width).reshape(
+                        packed, queries, width
+                    ),
+                    negative_inf,
+                )
+            return scores + mask_chunk.float().expand(
+                batch, query_heads, queries, width
+            ).reshape(packed, queries, width)
+
+        for start in range(0, key_len, chunk):
+            stop = min(start + chunk, key_len)
+            scores = _apply_mask(
+                torch.bmm(query_flat, _pack_kv(key_fp32, start, stop).transpose(1, 2)) * scale,
+                start,
+                stop,
+            )
+            block_max = scores.amax(dim=-1, keepdim=True)
+            new_max = torch.maximum(running_max, block_max)
+            running_sum = running_sum * torch.exp(running_max - new_max) + torch.exp(
+                scores - new_max
+            ).sum(dim=-1, keepdim=True)
+            running_max = new_max
+        grad_query_flat = torch.zeros_like(query_flat)
+        grad_key = torch.zeros_like(key_fp32)
+        grad_value = torch.zeros_like(value_fp32)
+        for start in range(0, key_len, chunk):
+            stop = min(start + chunk, key_len)
+            width = stop - start
+            key_flat = _pack_kv(key_fp32, start, stop)
+            value_flat = _pack_kv(value_fp32, start, stop)
+            scores = _apply_mask(
+                torch.bmm(query_flat, key_flat.transpose(1, 2)) * scale,
+                start,
+                stop,
+            )
+            probs = torch.exp(scores - running_max) / running_sum.clamp_min(
+                torch.finfo(running_sum.dtype).tiny
+            )
+            grad_value_packed = torch.bmm(probs.transpose(1, 2), grad_out_flat)
+            grad_value[:, :, start:stop, :] = grad_value_packed.view(
+                batch, key_heads, repeats, width, head_dim
+            ).sum(dim=2)
+            grad_probs = torch.bmm(grad_out_flat, value_flat.transpose(1, 2))
+            delta = grad_probs - (grad_probs * probs).sum(dim=-1, keepdim=True)
+            grad_scores = probs * delta
+            grad_query_flat = grad_query_flat + torch.bmm(grad_scores, key_flat) * scale
+            grad_key_packed = torch.bmm(grad_scores.transpose(1, 2), query_flat) * scale
+            grad_key[:, :, start:stop, :] = grad_key_packed.view(
+                batch, key_heads, repeats, width, head_dim
+            ).sum(dim=2)
+        return (
+            grad_query_flat.view(batch, query_heads, queries, head_dim).to(dtype=query.dtype),
+            grad_key.to(dtype=key.dtype),
+            grad_value.to(dtype=value.dtype),
+            None,
+            None,
+            None,
+        )
+
+
+def mps_safe_sdpa_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    dropout: float = 0.0,
+    scaling: float | None = None,
+    is_causal: bool | None = None,
+    position_bias: torch.Tensor | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, None]:
+    """Use compact GQA attention for long MPS decoder backwards; else official SDPA."""
+
+    kv_length = int(key.shape[-2])
+    use_chunked = (
+        query.device.type == "mps"
+        and torch.is_grad_enabled()
+        and query.requires_grad
+        and kv_length >= _LONG_KV_THRESHOLD
+        and dropout == 0.0
+        and position_bias is None
+        and not kwargs.get("output_attentions", False)
+    )
+    if not use_chunked:
+        return _original_sdpa_attention_forward(
+            module, query, key, value, attention_mask,
+            dropout=dropout, scaling=scaling, is_causal=is_causal,
+            position_bias=position_bias, **kwargs,
+        )
+    scale = 1.0 if scaling is None else float(scaling)
+    output = _MpsGqaChunkedAttention.apply(
+        query, key, value, attention_mask, scale, _LONG_KV_CHUNK,
+    )
+    return output.transpose(1, 2).contiguous(), None
+
+
+def register_mps_backends() -> None:
+    """Register the schema25 expert and attention implementations explicitly."""
+
+    ALL_EXPERTS_FUNCTIONS.register("lora_eager", lora_aware_eager_experts_forward)
+    ALL_EXPERTS_FUNCTIONS.register("modilify_mk2_lora", inference_experts_forward)
+    ALL_EXPERTS_FUNCTIONS.register("mps_grouped_mm", mps_grouped_mm_experts_forward)
+    ALL_EXPERTS_FUNCTIONS.register("mps_segmented", mps_segmented_experts_forward)
+    ALL_ATTENTION_FUNCTIONS.register("sdpa", mps_safe_sdpa_attention_forward)
+
+
+__all__ = [
+    "attach_expert_lora",
+    "grouped_expert_offsets",
+    "lora_aware_eager_experts_forward",
+    "mps_grouped_mm_experts_forward",
+    "mps_safe_sdpa_attention_forward",
+    "mps_segmented_experts_forward",
+    "register_mps_backends",
+]
+
+
+def inference_experts_forward(self, hidden_states, top_k_index, top_k_weights):
+    """Use the reference MPS path locally and the Torch fallback elsewhere."""
+    implementation = (mps_segmented_experts_forward if hidden_states.device.type == "mps"
+                      else lora_aware_eager_experts_forward)
+    return implementation(self, hidden_states, top_k_index, top_k_weights)

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.10
+transformers==5.14.1
+accelerate>=1.14
+safetensors>=0.8

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,96 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "local_files_only": true,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "processor_class": "Gemma4Processor",
+  "response_schema": {
+    "properties": {
+      "content": {
+        "type": "string"
+      },
+      "role": {
+        "const": "assistant"
+      },
+      "thinking": {
+        "type": "string"
+      },
+      "tool_calls": {
+        "items": {
+          "properties": {
+            "function": {
+              "properties": {
+                "arguments": {
+                  "additionalProperties": {},
+                  "type": "object",
+                  "x-parser": "gemma4-tool-call"
+                },
+                "name": {
+                  "type": "string"
+                }
+              },
+              "type": "object",
+              "x-regex": "call\\:(?P<name>\\w+)(?P<arguments>\\{.*\\})"
+            },
+            "type": {
+              "const": "function"
+            }
+          },
+          "type": "object"
+        },
+        "type": "array",
+        "x-regex-iterator": "<\\|tool_call>(.*?)<tool_call\\|>"
+      }
+    },
+    "type": "object",
+    "x-regex": "(\\<\\|channel\\>thought\\n(?P<thinking>.*?)\\<channel\\|\\>)?(?P<tool_calls>\\<\\|tool_call\\>.*\\<tool_call\\|\\>)?(?P<content>(?:(?!\\<turn\\|\\>)(?!\\<\\|tool_response\\>).)+)?(?:\\<turn\\|\\>|\\<\\|tool_response\\>)?"
+  },
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}

vocab_ops.py

@@ -0,0 +1,643 @@
+"""Exact memory-bounded vocabulary projection, sampling, and backward."""
+
+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."""
+
+    best, index = values.max(dim=-1)
+    width = values.shape[-1]
+    if width <= 0:
+        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,
+        top_k: int | None = 40,
+        min_p: float | None = 0.05,
+    ) -> 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]
+
+        use_constrained = top_k is not None and top_k > 0
+        chunk_cand_scores = []
+        chunk_cand_tokens = []
+
+        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,
+                )
+
+                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:
+                    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
+
+        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
+            uniform = torch.rand(
+                eligible.shape,
+                device=eligible.device,
+                dtype=torch.float32,
+            ).clamp_(
+                min=torch.finfo(torch.float32).tiny,
+                max=1.0 - torch.finfo(torch.float32).eps,
+            )
+            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)
+
+        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,
+            sample_log_z - selected_score,
+        )
+
+    @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,
+        grad_sampled_nll: 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
+        )
+        if grad_sampled_nll is not None:
+            confidence_scale = confidence_scale - grad_sampled_nll.float() / ctx.temperature
+        hidden_fp32 = hidden.float()
+        for start in range(0, weight.shape[0], ctx.chunk_size):
+            stop = min(start + ctx.chunk_size, weight.shape[0])
+            # The B16 training shape is 4,096 rows by a 32,768-wide vocabulary
+            # chunk.  MPS BF16 GEMM can return NaNs for this backward-only
+            # projection even when every saved activation is finite.  Recompute
+            # scores and accumulate the exact VJP in FP32; cast the finished
+            # hidden gradient back to the model dtype only once below.
+            weight_fp32 = weight[start:stop].float()
+            scores = torch.tanh(
+                F.linear(hidden_fp32, weight_fp32) / ctx.softcap
+            ) * ctx.softcap
+            # Non-deterministic MPS GEMM can make a recomputed score slightly
+            # exceed the saved log-partition.  Clamp the exponent so softmax
+            # probabilities stay in [0, 1] instead of overflowing to Inf.
+            raw_probability = torch.exp(
+                (scores - raw_log_z[:, None]).clamp(max=0.0)
+            )
+            sample_logit = (
+                scores / ctx.temperature - sample_log_z[:, None]
+            ).clamp(max=0.0)
+            sample_probability = torch.exp(sample_logit)
+            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()
+                entropy_grad = -(
+                    grad_entropy_f / ctx.temperature
+                )[:, None] * sample_probability * (
+                    sample_logit + 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()
+            score_gradient = torch.nan_to_num(
+                score_gradient, nan=0.0, posinf=0.0, neginf=0.0
+            )
+            chunk_gradient = F.linear(
+                score_gradient,
+                weight_fp32.transpose(0, 1),
+            )
+            grad_hidden.add_(chunk_gradient)
+
+        grad_hidden = torch.nan_to_num(
+            grad_hidden, nan=0.0, posinf=0.0, neginf=0.0
+        )
+        return (
+            grad_hidden.to(hidden.dtype),
+            None,
+            None,
+            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,
+    top_k: int | None = 40,
+    min_p: float | None = 0.05,
+    return_sampled_nll: bool = False,
+) -> 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),
+        top_k,
+        min_p,
+    )
+    (
+        raw_nll,
+        temperature_nll,
+        proposal,
+        confidence,
+        entropy,
+        greedy,
+        greedy_confidence,
+        sampled_nll,
+    ) = outputs
+    result = (
+        raw_nll.view(shape),
+        temperature_nll.view(shape),
+        proposal.view(shape),
+        confidence.view(shape),
+        entropy.view(shape),
+        greedy.view(shape),
+        greedy_confidence.view(shape),
+    )
+    return result + (sampled_nll.view(shape),) if return_sampled_nll else result
+
+
+@torch.no_grad()
+def chunked_vocab_statistics(
+    hidden: torch.Tensor,
+    weight: torch.Tensor,
+    *,
+    softcap: float,
+    temperature: float,
+    chunk_size: int,
+    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,
+    torch.Tensor,
+    torch.LongTensor,
+    torch.Tensor,
+]:
+    """Return exact inference statistics, optionally penalizing seen tokens."""
+
+    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]:
+            raise ValueError(
+                "`sampling_generators` must contain one generator per batch row."
+            )
+
+        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:
+                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)
+            )
+
+            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,
+                        device=sample_scores.device,
+                        dtype=torch.float32,
+                    )
+                else:
+                    per_request_shape = (
+                        rows_per_batch,
+                        sample_scores.shape[-1],
+                    )
+                    uniform = torch.cat(
+                        [
+                            torch.rand(
+                                per_request_shape,
+                                device=sample_scores.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,
+                )
+                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
+
+        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,
+            )
+            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
+        )
+        return (
+            selected.view(output_shape),
+            confidence.view(output_shape),
+            entropy.view(output_shape),
+            greedy.view(output_shape),
+            greedy_confidence.view(output_shape),
+        )
+
+    labels = torch.full(
+        hidden.shape[:-1],
+        -100,
+        device=hidden.device,
+        dtype=torch.long,
+    )
+    (
+        _,
+        _,
+        proposal,
+        confidence,
+        entropy,
+        greedy,
+        greedy_confidence,
+    ) = chunked_vocab_training_statistics(
+        hidden,
+        weight,
+        labels,
+        softcap=softcap,
+        temperature=temperature,
+        chunk_size=chunk_size,
+        top_k=top_k,
+        min_p=min_p,
+    )
+    return proposal, confidence, entropy, greedy, greedy_confidence
+
+
+__all__ = [
+    "chunked_vocab_statistics",
+    "chunked_vocab_training_statistics",
+]