Publish Modilify Mk2 Preview MLX

.gitattributes

@@ -1,35 +1,3 @@
-*.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
+assets/01-LOGO.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,11 @@
+.DS_Store
+__pycache__/
+.vscode/
+*.py[cod]
+.venv/
+build/
+dist/
+*.egg-info/
+.cache/
+.env
+/.upload-fast.sh

LICENSE

@@ -0,0 +1,125 @@
+Modilify Open Model License 1.0
+
+Copyright 2026 Modilify
+
+1. Definitions
+
+"Model" means the weights, configuration, inference code, tokenizer, processor,
+and documentation distributed with this license. "Derivative Model" means a
+modified, fine-tuned, distilled, merged, quantized, or otherwise adapted version
+of the Model. "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 this license, Modilify grants You a worldwide, perpetual,
+non-exclusive, royalty-free, irrevocable copyright license to use, reproduce,
+prepare derivative works of, publicly display, publicly perform, sublicense,
+host as a service, and distribute the Model and Derivative Models, including for
+commercial purposes.
+
+3. 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, offer to sell, sell, import, and otherwise transfer the Model 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.
+
+4. Conditions on Redistribution
+
+If You distribute the Model or a Derivative Model, You must:
+
+a. provide recipients a copy of this license;
+b. retain copyright, patent, attribution, and NOTICE statements;
+c. state clearly that You modified the Model and identify material modifications;
+d. preserve applicable third-party license and attribution notices; and
+e. publish with the distributed model a reasonably accessible impact statement
+   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.
+
+5. 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.
+
+6. 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.
+
+7. Third-Party Components
+
+The Model includes or derives from third-party components identified in
+NOTICE.md. Those components remain subject to their applicable
+licenses and terms. In particular, rights and obligations associated with the
+Google DiffusionGemma base are not removed, narrowed, or replaced by this
+license. You are responsible for complying with all applicable upstream terms.
+
+8. 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.
+
+9. 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.
+
+10. 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.
+
+11. Governing Law and Venue
+
+This license is governed by the laws of the State of California, excluding its
+conflict-of-law rules. Any dispute arising from this license must be brought in
+the state or federal courts located in Santa Clara County, California, and each
+party consents to their personal jurisdiction and venue.
+
+12. Entire License; Severability
+
+This document states the complete Modilify license for the Model, subject to
+applicable third-party terms. If a provision is unenforceable, it will be limited
+to the minimum extent necessary and the remaining provisions remain effective.

NOTICE.md

@@ -0,0 +1,255 @@
+# Modilify Mk2 — Notices
+
+Copyright 2026 Modilify
+
+This distribution is derived from `google/diffusiongemma-26B-A4B-it`, published by Google DeepMind under Apache License 2.0. This MLX distribution retains the text encoder/decoder backbone, tokenizer, and chat template. It does not include vision tower, vision projection, or image/video inference path. Modilify added trained LoRA adapters, dual-timescale GDN2 trajectory memory, working and persistent memory readers, and the confidence-and-entropy prefix commit policy. The native MLX implementation uses the DiffusionGemma text architecture from mlx-vlm and unfused adapters from mlx-lm.
+
+The Modilify Open Model License 1.0 applies to Modilify's distribution and original contributions. It does not erase, narrow, or replace rights and notices applicable to upstream components. Users remain responsible for complying with all applicable upstream terms.
+
+- Upstream model: https://huggingface.co/google/diffusiongemma-26B-A4B-it
+- Transformers project: https://github.com/huggingface/transformers
+- MLX project: https://github.com/ml-explore/mlx
+- mlx-vlm project: https://github.com/Blaizzy/mlx-vlm
+- mlx-lm project: https://github.com/ml-explore/mlx-lm
+
+The configuration subclasses public DiffusionGemma interfaces in Hugging Face Transformers, distributed under Apache License 2.0. MLX and mlx-lm are distributed under the MIT License, and mlx-vlm under the MIT License. These libraries are installed as dependencies; their original notices remain in their distributions.
+
+## Derivative Model Impact Statement Template
+
+When distributing a derivative of Modilify Mk2, include a public impact statement covering the following items. No separate submission to Modilify is required.
+
+### Identity and modifications
+
+- Model name, version, publisher, and contact.
+- Base version.
+- Material modifications, data sources, merges, quantization, or adaptation.
+
+### Intended and excluded uses
+
+- Intended users and use cases.
+- Explicitly excluded uses.
+- Deployment context and degree of human oversight.
+
+### Evaluation scope
+
+- Evaluated capabilities and datasets.
+- Languages, modalities, populations, or contexts not evaluated.
+- Hardware and software used.
+
+### Known limitations and foreseeable risks
+
+- Reliability limitations.
+- Safety, bias, privacy, security, and misuse risks.
+- High-risk decisions the model must not make autonomously.
+
+### Mitigations and monitoring
+
+- Technical and organizational safeguards.
+- Human review, appeal, and correction mechanisms.
+- Monitoring, incident response, and update policy.
+
+## Apache License 2.0
+
+The complete license text applicable to the upstream components follows.
+
+                                 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.

README.md

@@ -0,0 +1,126 @@
+---
+license: other
+license_name: modilify-open-model-license-1.0
+license_link: LICENSE
+library_name: mlx
+pipeline_tag: text-generation
+tags:
+  - mlx
+  - diffusion
+  - mixture-of-experts
+  - custom-code
+  - modilify-mk2
+---
+
+![Modilify](assets/01-LOGO.jpg)
+
+# Modilify Mk2 Preview · MLX
+
+Modilify Mk2 is a rolling-canvas text diffusion model with learned trajectory memory. This release provides complete, unquantized native MLX weights and an inference package for Apple Silicon Macs. It builds on [Google DiffusionGemma](https://huggingface.co/google/diffusiongemma-26B-A4B-it), adding LoRA adapters, GDN2 memory at two timescales, and a confidence-and-entropy prefix commit policy.
+
+The model revises a 256-token canvas over successive denoising passes, commits a variable-length prefix, and carries memory forward as the canvas advances. Per-position and row-level working states update on every denoising pass; persistent memory updates only when tokens commit.
+
+This is an experimental **text-only preview**. The complete model is included; no separate base-model download is required. PyTorch is not required.
+
+## Download and run
+
+Use Python 3.12 on an Apple Silicon Mac with macOS 26 or newer for the pinned MLX environment. The weights occupy **48.23 GiB**. Inference needs additional memory for KV caches, recurrent states, and computation. See [local comparison](#local-comparison) for measured short-request memory usage.
+
+```bash
+python3.12 -m venv .venv
+source .venv/bin/activate
+python -m pip install "huggingface-hub==1.20.1"
+
+hf download Modilify/Modilify-Mk2-preview-mlx --local-dir ./Modilify-Mk2-preview-mlx
+cd Modilify-Mk2-preview-mlx
+python -m pip install .
+
+python inference.py --prompt "Why is the sky blue?" \
+  --max-new-tokens 256 --max-denoising-steps 256 --stream
+```
+
+The source entrypoint defaults to the model in the same directory. The installed command accepts a local model directory or a Hub model ID:
+
+```bash
+modilify-mlx --model ./ --prompt "你好，请介绍一下自己。" --stream
+modilify-mlx --model Modilify/Modilify-Mk2-preview-mlx \
+  --prompt "What is 2 + 3?" --max-new-tokens 128 --max-denoising-steps 128
+```
+
+The Hub ID path downloads the model snapshot into the Hugging Face cache. Later runs reuse the cached files. All inference computation runs locally.
+
+Thinking is enabled by default; use `--think False` to disable it. `--max-new-tokens` limits output length; `--max-denoising-steps` limits denoising work. Reaching either budget may truncate a response. The CLI defaults to 8192 output tokens and the generation configuration's 48 denoising steps. Set both budgets explicitly for your workload. `--canvas-length` can reduce the canvas from its maximum of 256. Without `--stream`, the CLI emits JSONL events and a final result with text, token IDs, stop reason, and timing/memory metrics.
+
+## Python API
+
+Load once and reuse the runtime for independent requests:
+
+```python
+from inference import load_model, generate
+
+runtime = load_model("./")
+result = generate(
+    runtime,
+    "Explain diffusion models in one paragraph.",
+    max_new_tokens=256,
+    max_denoising_steps=256,
+    seed=42,
+    think=True,
+)
+print(result["text"])
+print(result["stop_reason"])
+```
+
+Pass `messages=[{"role": "user", "content": "..."}]` instead of `prompt` for a conversation. The tokenizer's Modilify template handles system, user, and assistant turns. With thinking enabled, generated text may include reasoning before the answer.
+
+## Continuous requests
+
+```bash
+printf '%s\n' \
+  '{"request_id":"a","prompt":"你好","max_new_tokens":128,"max_denoising_steps":128}' \
+  '{"request_id":"b","prompt":"What is 2 + 3?","seed":42,"think":false,"max_new_tokens":128,"max_denoising_steps":128}' | \
+  modilify-mlx --model ./ --requests-jsonl - \
+    --batch-size 4 --pipeline-depth 2 --max-batch-tokens 1024
+```
+
+Each line requires a unique `request_id` and either `prompt` or `messages`. Optional fields are `max_new_tokens`, `max_denoising_steps`, `seed`, and `think`. The default scheduler interleaves independent rows and keeps their random streams and states separate. `--batch-size` controls resident requests; `--pipeline-depth` bounds in-flight denoises.
+
+## Release details
+
+| Item | Value |
+| --- | --- |
+| Base model | DiffusionGemma 26B-A4B-it text backbone |
+| Saved training step | 1250 |
+| Memory topology | schema25 / `compact_gdn2_v2` |
+| Memory scheme | `dual_timescale_gdn2_trajectory_memory` |
+| Canvas | Up to 256 tokens |
+| Expert routing | 8 of 128 experts per token |
+| Dense / expert LoRA rank | 16 / 8 |
+| Weight precision | BF16, with small trajectory norms and dynamics parameters in FP32 |
+| Proposal constraints | top_k 40, min_p 0.05 |
+| Commit defaults | failure budget 0.2, target confidence 0.5 |
+| Weights | 32 Safetensors shards, 1323 tensors |
+
+The complete backbone, adapters, and memory parameters are stored in the native MLX parameter layout. Adapters remain unfused to preserve the reference BF16 computation. Loading validates model topology, tensor shape/dtype, and shard SHA256 hashes. `export_manifest.json` contains the weight index metadata and release provenance; no optimizer state or training data is shipped.
+
+This model uses its own rolling-canvas generation loop. Use the included `inference.py`, `modilify-mlx`, or Python API. Generic `mlx_lm.generate`, Transformers `AutoModel.from_pretrained`, and standard autoregressive serving backends do not implement this protocol. Image, audio, and video inputs are not supported by this release.
+
+## Training, evaluation, and limitations
+
+The text backbone is adapted with dense and expert LoRA and trained trajectory memory. Adaptation uses response supervision over the valid rolling canvas, with terminal, confidence, and commit-readiness objectives. This release uses the saved step 1250 weights. The adaptation dataset is not distributed and its source composition is not documented in this release; refer to the upstream model card for base-model training information.
+
+Local comparisons establish export correctness for the tested cases, not task quality or safety. No capability benchmark score is claimed for this preview, and upstream benchmark scores do not describe this modified model. Long-context quality, broad language coverage, and deployment safety have not been established by those checks.
+
+## Local comparison
+
+On 2026-10-06, the package installed in a fresh Python 3.12 environment without PyTorch and ran offline. Two requests matched the strict step 1250 source checkpoint at 1/2/5/10/20 denoising steps, including token IDs, stopping and canvas shifts. At five steps, interleaved B2 also matched independent B1 outputs and all three GDN2 state hashes. Short-request MLX peak allocations were about 50.79 GiB; longer workloads require additional memory.
+
+## Intended use and impact statement
+
+This preview is intended for research on text diffusion, trajectory memory, local generation, and reproducible inference experiments. Generated statements can be incorrect, biased, or inappropriate; verify consequential outputs and keep qualified human review for sensitive decisions. The model is not intended to make autonomous high-risk decisions. Do not assume local execution alone provides safety, privacy, or resistance to prompt injection. Deployment should include use-specific testing, appropriate safeguards, and incident monitoring.
+
+## License and attribution
+
+Modilify's original contributions and distribution are covered by the [Modilify Open Model License 1.0](LICENSE), a custom license. Upstream components retain their own terms; the DiffusionGemma base is published under Apache 2.0. See [NOTICE.md](NOTICE.md) for attribution and upstream notices.
+
+Maintained by **Modilify**. Report model and runtime issues through the [Hub discussions](https://huggingface.co/Modilify/Modilify-Mk2-preview-mlx/discussions).

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 -%}

config.json

@@ -0,0 +1,134 @@
+{
+  "canvas_length": 256,
+  "channel_end_token_id": 101,
+  "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,
+  "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,
+  "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"
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "eos_token_id": 1,
+  "jump_on_no_progress_after": 12,
+  "max_denoising_steps": null,
+  "max_ponder_steps": 64,
+  "min_trajectory_progress": 0.005,
+  "repetition_penalty": 1.0,
+  "repetition_penalty_exclude_token_ids": [],
+  "turn_end_token_id": 106
+}

inference.py

@@ -0,0 +1,1186 @@
+"""Native MLX schema25 inference with bounded KV reuse and continuous admission."""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import hashlib
+import json
+import math
+import queue
+import sys
+import threading
+import time
+from collections import OrderedDict, deque
+from dataclasses import dataclass, fields, replace
+from pathlib import Path
+from typing import Any, Iterator
+
+import mlx.core as mx
+
+from modilify_mk2.configuration_modilify_mk2 import DENOISE_TEMPERATURE
+from modilify_mk2.mlx_commit_policy import (
+    fused_commit_failure_rate, infer_commit_reason, select_commit_lengths,
+)
+from modilify_mk2.mlx_model import MLXCanvasOutput, MLXModilifyMk2
+from modilify_mk2.runtime import MLXRuntime, load_runtime
+from modilify_mk2.mlx_state import MLXLatentState, MLXRollingState
+from modilify_mk2.chat import apply_chat_template
+def parse_bool(value: str | bool) -> bool:
+    """Parse explicit CLI booleans such as ``--think true``."""
+
+    if isinstance(value, bool):
+        return value
+    normalized = value.strip().lower()
+    if normalized in {"1", "true", "yes", "on"}:
+        return True
+    if normalized in {"0", "false", "no", "off"}:
+        return False
+    raise argparse.ArgumentTypeError("expected true or false")
+
+
+REQUEST_FIELDS = frozenset(
+    {
+        "request_id",
+        "prompt",
+        "messages",
+        "max_new_tokens",
+        "max_denoising_steps",
+        "seed",
+        "think",
+    }
+)
+
+
+@dataclass(frozen=True)
+class ContinuousRequest:
+    """One validated machine-mode request before chat-template tokenization."""
+
+    request_id: str
+    messages: list[dict[str, Any]]
+    max_new_tokens: int
+    max_denoising_steps: int | None
+    seed: int
+    think: bool
+    prompt: str | None = None
+
+
+def stable_request_seed(base_seed: int, request_id: str) -> int:
+    """Derive a scheduler-independent non-negative seed from a stable request ID."""
+
+    payload = f"{base_seed}\0{request_id}".encode("utf-8")
+    return int.from_bytes(hashlib.sha256(payload).digest()[:8], "big") & ((1 << 63) - 1)
+
+
+def _positive_int(record: dict[str, Any], name: str, default: int | None) -> int | None:
+    value = record.get(name, default)
+    if value is None:
+        return default
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(f"`{name}` must be a positive integer or null.")
+    return value
+
+
+def parse_continuous_request(
+    record: Any,
+    *,
+    default_max_new_tokens: int,
+    default_max_denoising_steps: int | None,
+    default_seed: int,
+    default_think: bool,
+) -> ContinuousRequest:
+    """Validate one strict request object without silently inventing identity."""
+
+    if not isinstance(record, dict):
+        raise ValueError("Each request line must be a JSON object.")
+    unknown = sorted(set(record).difference(REQUEST_FIELDS))
+    if unknown:
+        raise ValueError(f"Unknown request fields: {unknown}")
+    request_id = record.get("request_id")
+    if not isinstance(request_id, str) or not request_id.strip():
+        raise ValueError("`request_id` must be a non-empty string.")
+    has_prompt = "prompt" in record
+    has_messages = "messages" in record
+    if has_prompt == has_messages:
+        raise ValueError("Exactly one of `prompt` or `messages` is required.")
+
+    prompt = None
+    if has_prompt:
+        prompt = record["prompt"]
+        if not isinstance(prompt, str):
+            raise ValueError("`prompt` must be a string.")
+        messages = [{"role": "user", "content": prompt}]
+    else:
+        messages = record["messages"]
+        if not isinstance(messages, list) or not messages:
+            raise ValueError("`messages` must be a non-empty list.")
+        if any(
+            not isinstance(message, dict)
+            or "role" not in message
+            or "content" not in message
+            for message in messages
+        ):
+            raise ValueError("Every message must be an object with `role` and `content`.")
+
+    think = record.get("think", default_think)
+    if not isinstance(think, bool):
+        raise ValueError("`think` must be a boolean.")
+    supplied_seed = record.get("seed")
+    if supplied_seed is not None and (
+        isinstance(supplied_seed, bool) or not isinstance(supplied_seed, int)
+    ):
+        raise ValueError("`seed` must be an integer or null.")
+    seed = (
+        stable_request_seed(default_seed, request_id)
+        if supplied_seed is None
+        else supplied_seed
+    )
+    return ContinuousRequest(
+        request_id=request_id,
+        prompt=prompt,
+        messages=messages,
+        max_new_tokens=int(
+            _positive_int(record, "max_new_tokens", default_max_new_tokens)
+        ),
+        max_denoising_steps=_positive_int(
+            record, "max_denoising_steps", default_max_denoising_steps
+        ),
+        seed=seed,
+        think=think,
+    )
+
+
+def _extract_input_ids(encoded: Any) -> list[int]:
+    value = encoded.get("input_ids") if isinstance(encoded, dict) else encoded
+    if hasattr(encoded, "input_ids"):
+        value = encoded.input_ids
+    if isinstance(value, tuple):
+        value = list(value)
+    if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list):
+        value = value[0]
+    if not isinstance(value, list) or any(
+        isinstance(token, bool) or not isinstance(token, int) for token in value
+    ):
+        raise ValueError("Chat template did not return one integer token sequence.")
+    if not value:
+        raise ValueError("Chat template returned an empty prompt.")
+    return [int(token) for token in value]
+
+
+def _seed(value: int, salt: int) -> mx.array:
+    return mx.random.key((int(value) + int(salt)) % (2**32))
+
+
+def _logical(value: mx.array, head: mx.array) -> mx.array:
+    canvas = value.shape[1]
+    index = (head[:, None] + mx.arange(canvas)[None, :]) % canvas
+    return mx.take_along_axis(value, index, axis=1)
+
+
+def _physical(value: mx.array, head: mx.array) -> mx.array:
+    canvas = value.shape[1]
+    index = (mx.arange(canvas)[None, :] - head[:, None]) % canvas
+    return mx.take_along_axis(value, index, axis=1)
+
+
+def _concat_rows(values: list[Any]) -> Any:
+    first = values[0]
+    if isinstance(first, mx.array):
+        return mx.concatenate(values, axis=0)
+    return type(first)(**{
+        item.name: _concat_rows([getattr(value, item.name) for value in values])
+        for item in fields(first)
+    })
+
+
+def _slice_row(value: Any, row: int) -> Any:
+    if isinstance(value, mx.array):
+        return value[row:row + 1]
+    return type(value)(**{
+        item.name: _slice_row(getattr(value, item.name), row)
+        for item in fields(value)
+    })
+
+
+def _arrays(value: Any) -> list[mx.array]:
+    if isinstance(value, mx.array):
+        return [value]
+    return [array for item in fields(value)
+            for array in _arrays(getattr(value, item.name))]
+
+
+def _clone_cache(cache: list[Any], *, compact: bool = False) -> list[Any]:
+    copied = []
+    arrays = []
+    for layer in cache:
+        clone = copy.copy(layer)
+        source = layer.state if compact else (layer.keys, layer.values)
+        if source[0] is not None:
+            clone.keys = mx.array(source[0])
+            clone.values = mx.array(source[1])
+            arrays.extend((clone.keys, clone.values))
+        copied.append(clone)
+    if arrays:
+        mx.eval(*arrays)
+    return copied
+
+
+@dataclass
+class _PrefixEntry:
+    cache: list[Any]
+    bytes: int
+
+
+class PrefixKVCache:
+    """LRU cache of immutable, block-boundary prompt KV snapshots."""
+
+    def __init__(self, limit_bytes: int):
+        self.limit_bytes = limit_bytes
+        self.entries: OrderedDict[tuple[int, ...], _PrefixEntry] = OrderedDict()
+        self.bytes = 0
+        self.hits = 0
+        self.reused_tokens = 0
+
+    def longest(self, ids: tuple[int, ...]) -> tuple[int, list[Any] | None]:
+        lengths = sorted(
+            {len(key) for key in self.entries if len(key) <= len(ids)},
+            reverse=True,
+        )
+        for length in lengths:
+            key = ids[:length]
+            entry = self.entries.get(key)
+            if entry is not None:
+                self.entries.move_to_end(key)
+                self.hits += 1
+                self.reused_tokens += length
+                return length, _clone_cache(entry.cache)
+        return 0, None
+
+    def put(self, ids: tuple[int, ...], cache: list[Any]) -> None:
+        if self.limit_bytes <= 0 or ids in self.entries:
+            return
+        size = sum(int(value.nbytes) for layer in cache
+                   for value in (layer.state if layer.keys is not None else ())
+                   if value is not None)
+        if size > self.limit_bytes:
+            return
+        snapshot = _clone_cache(cache, compact=True)
+        while self.entries and self.bytes + size > self.limit_bytes:
+            _, victim = self.entries.popitem(last=False)
+            self.bytes -= victim.bytes
+        self.entries[ids] = _PrefixEntry(snapshot, size)
+        self.bytes += size
+
+
+@dataclass
+class InferenceRow:
+    request: ContinuousRequest
+    prompt_ids: tuple[int, ...]
+    cache: list[Any]
+    rolling: MLXRollingState
+    generated: list[int]
+    repetition_seen: set[int]
+    created: float
+    admitted: float
+    denoise_steps: int = 0
+    jumps: int = 0
+    shifts: int = 0
+    first_token_at: float | None = None
+    stop_reason: str | None = None
+    first_denoise_at: float | None = None
+    last_denoise_at: float | None = None
+    prefill_seconds: float = 0.0
+
+
+@dataclass
+class PrefillRow:
+    request: ContinuousRequest
+    prompt_ids: tuple[int, ...]
+    cache: list[Any]
+    offset: int
+    reused: int
+    created: float
+    seconds: float = 0.0
+
+
+@dataclass
+class _DenoiseWork:
+    rows: list[InferenceRow]
+    rolling: MLXRollingState
+    output: Any
+    proposal: mx.array
+    remaining: mx.array
+    physical_positions: mx.array
+    next_latent: MLXLatentState
+    policy: Any
+
+
+def _forward_independent_rows(model: MLXModilifyMk2,
+                              rows: list[InferenceRow]) -> list[MLXCanvasOutput]:
+    """Interleave decoder layers while preserving independent singleton rows."""
+    decoder = model.model.decoder
+    latent = model.latent_deliberation
+    working_bus, persistent_bus = latent.working_memory_bus, latent.persistent_memory_bus
+    contexts = []
+    for row in rows:
+        rolling = row.rolling
+        canvas, state, head = rolling.canvas, rolling.latent, rolling.head
+        batch, length = canvas.shape
+        tokens = decoder.embed_tokens(canvas) * decoder.embed_scale
+        working, next_state = latent(token_embeddings=tokens,
+                                     confidence=state.confidence,
+                                     entropy=state.entropy, state=state,
+                                     canvas_head=head)
+        physical = (head[:, None] + mx.arange(length)[None, :]) % length
+        gather = mx.broadcast_to(physical[:, :, None], tokens.shape)
+        logical_tokens = mx.take_along_axis(tokens, mx.stop_gradient(gather), axis=1)
+        logical_working = mx.take_along_axis(working, mx.stop_gradient(gather), axis=1)
+        hidden = model._merge_context(logical_tokens, logical_working)
+        prefix_length = int(row.cache[0].offset)
+        full_mask = mx.ones((batch, prefix_length + length), mx.bool_)
+        masks = decoder._make_decoder_masks(hidden, row.cache, full_mask)
+        seen = latent.logical_seen(state, head)
+        contexts.append({"hidden": hidden, "working": working, "state": next_state,
+                         "tokens": tokens, "head": head, "masks": masks,
+                         "offset": prefix_length,
+                         "working_kv": working_bus.prepare_kv((logical_working, seen)),
+                         "persistent_kv": persistent_bus.prepare_kv(state.memory_slots, seen)})
+    reader = 0
+    for index, layer in enumerate(decoder.layers):
+        for row, context in zip(rows, contexts, strict=True):
+            hidden = layer(context["hidden"], context["masks"][layer.layer_type],
+                           row.cache[index], decoder=True, offset=context["offset"])
+            if layer.layer_type == "full_attention":
+                if context["working_kv"] is not None and reader < working_bus.num_readers:
+                    hidden = working_bus.read(hidden, reader, *context["working_kv"])
+                if context["persistent_kv"] is not None and reader < persistent_bus.num_readers:
+                    hidden = persistent_bus.read(hidden, reader, *context["persistent_kv"])
+            context["hidden"] = hidden
+        # Explicit layer boundaries keep the execution order interleaved and
+        # bound intermediate activations to the configured pipeline depth.
+        mx.async_eval(*(context["hidden"] for context in contexts))
+        if layer.layer_type == "full_attention":
+            reader += 1
+    outputs = []
+    for context in contexts:
+        hidden = decoder.norm(context["hidden"])
+        length = hidden.shape[1]
+        inverse = (mx.arange(length)[None, :] - context["head"][:, None]) % length
+        heavy = mx.take_along_axis(hidden, mx.stop_gradient(mx.broadcast_to(
+            inverse[:, :, None], hidden.shape)), axis=1)
+        outputs.append(MLXCanvasOutput(heavy, context["working"],
+                                       context["state"], context["tokens"]))
+    return outputs
+
+
+def _noise(seed: int, step: int, canvas: int, vocab: int) -> mx.array:
+    return mx.random.randint(0, vocab, (1, canvas),
+                             key=_seed(seed, 0x51A7 + step * 1000003))
+
+
+def _empty_rolling(config: Any, seed: int, max_new_tokens: int,
+                   pad_token_id: int) -> MLXRollingState:
+    canvas = int(config.canvas_length)
+    vocab = int(config.text_config.vocab_size)
+    latent = MLXLatentState.empty(1, canvas, int(config.latent_memory_slots),
+                                  int(config.latent_dim), enable_gdn2=True)
+    latent = replace(latent, entropy=mx.full((1, canvas), math.log(vocab), mx.float32))
+    initial = mx.where(mx.arange(canvas)[None, :] < max_new_tokens,
+                       _noise(seed, 0, canvas, vocab), pad_token_id)
+    return MLXRollingState(initial, latent,
+                            mx.zeros((1,), mx.int32))
+
+
+def _inference_statistics(hidden: mx.array, weight: mx.array,
+                          rows: list[InferenceRow], *, softcap: float,
+                          chunk_size: int, penalty: float,
+                          excluded: set[int],
+                          top_k: int | None = 40,
+                          min_p: float | None = 0.05) -> tuple[mx.array, ...]:
+    """Exact chunked Top-K and Min-P Gumbel sampling without a retained full-vocabulary matrix."""
+    batch, canvas, dim = hidden.shape
+    flat = hidden.reshape(batch * canvas, dim)
+    neg_inf = mx.full((batch * canvas,), -mx.inf, mx.float32)
+    log_z = neg_inf
+    best_gumbel = neg_inf
+    chosen_score = mx.zeros_like(log_z)
+    chosen = mx.zeros((batch * canvas,), mx.int32)
+    greedy_score = neg_inf
+    greedy = mx.zeros_like(chosen)
+    moment_max = neg_inf
+    moment_sum = mx.zeros_like(log_z)
+    moment_weighted = mx.zeros_like(log_z)
+    repetition_mask = None
+    if penalty != 1.0:
+        masks = []
+        for row in rows:
+            mask = mx.zeros((weight.shape[0],), mx.bool_)
+            eligible = sorted(row.repetition_seen - excluded)
+            if eligible:
+                mask[mx.array(eligible, mx.int32)] = True
+            masks.append(mask)
+        repetition_mask = mx.stack(masks, axis=0)
+
+    use_constrained = top_k is not None and top_k > 0
+    chunk_cand_scores = []
+    chunk_cand_tokens = []
+
+    for chunk_number, start in enumerate(range(0, weight.shape[0], chunk_size)):
+        stop = min(start + chunk_size, weight.shape[0])
+        score = mx.tanh((flat @ weight[start:stop].T).astype(mx.float32) / softcap) * softcap
+        if penalty != 1.0:
+            assert repetition_mask is not None
+            row_scores = score.reshape(batch, canvas, stop - start)
+            seen = repetition_mask[:, None, start:stop]
+            changed = mx.where(row_scores < 0, row_scores * penalty, row_scores / penalty)
+            score = mx.where(seen, changed, row_scores).reshape(batch * canvas, stop - start)
+        score = score / DENOISE_TEMPERATURE
+        log_z = mx.logaddexp(log_z, mx.logsumexp(score, axis=-1))
+
+        local_max = mx.max(score, axis=-1)
+        better_greedy = local_max > greedy_score
+        greedy = mx.where(better_greedy, mx.argmax(score, axis=-1).astype(mx.int32) + start, greedy)
+        greedy_score = mx.maximum(greedy_score, local_max)
+
+        shifted = mx.exp(score - local_max[:, None])
+        next_max = mx.maximum(moment_max, local_max)
+        old_scale = mx.exp(moment_max - next_max)
+        new_scale = mx.exp(local_max - next_max)
+        moment_sum = moment_sum * old_scale + mx.sum(shifted, axis=-1) * new_scale
+        moment_weighted = moment_weighted * old_scale + mx.sum(shifted * score, axis=-1) * new_scale
+        moment_max = next_max
+
+        if use_constrained:
+            k = min(top_k, stop - start)
+            chunk_top_idx = mx.stop_gradient(mx.argpartition(-score, kth=k - 1, axis=-1)[:, :k]).astype(mx.int32)
+            chunk_top_score = mx.take_along_axis(score, chunk_top_idx, axis=-1)
+            chunk_cand_scores.append(chunk_top_score)
+            chunk_cand_tokens.append(chunk_top_idx + start)
+        else:
+            uniform = mx.concatenate([
+                mx.random.uniform(shape=(canvas, stop - start),
+                                  key=_seed(row.request.seed,
+                                            0xC09A + row.denoise_steps * 1000003 + chunk_number))
+                for row in rows
+            ], axis=0)
+            uniform = mx.clip(uniform, 1.17549435e-38, 1.0 - 1.19209290e-7)
+            gumbel = score - mx.log(-mx.log(uniform))
+            local_index = mx.argmax(gumbel, axis=-1)
+            local_best = mx.max(gumbel, axis=-1)
+            better = local_best > best_gumbel
+            chosen_score = mx.where(better, mx.take_along_axis(score, local_index[:, None], axis=-1)[:, 0], chosen_score)
+            chosen = mx.where(better, local_index + start, chosen)
+            best_gumbel = mx.maximum(best_gumbel, local_best)
+
+        # Materialize each chunk without a CPU/GPU barrier. Include candidates
+        # so their argpartition graph does not retain every vocabulary slab.
+        live = [log_z, greedy, greedy_score, moment_max, moment_sum, moment_weighted]
+        if use_constrained:
+            live.extend((chunk_cand_scores[-1], chunk_cand_tokens[-1]))
+        else:
+            live.extend((best_gumbel, chosen_score, chosen))
+        mx.async_eval(*live)
+
+    if use_constrained:
+        all_cand_scores = mx.concatenate(chunk_cand_scores, axis=-1)
+        all_cand_tokens = mx.concatenate(chunk_cand_tokens, axis=-1)
+        global_k = min(top_k, all_cand_scores.shape[-1])
+        global_idx = mx.stop_gradient(mx.argpartition(-all_cand_scores, kth=global_k - 1, axis=-1)[:, :global_k]).astype(mx.int32)
+        cand_scores = mx.take_along_axis(all_cand_scores, global_idx, axis=-1)
+        cand_tokens = mx.take_along_axis(all_cand_tokens, global_idx, axis=-1)
+        if min_p is not None and min_p > 0.0:
+            thresh = greedy_score + math.log(float(min_p))
+            valid_cand = cand_scores >= thresh[:, None]
+            eligible_scores = mx.where(valid_cand, cand_scores, -mx.inf)
+        else:
+            eligible_scores = cand_scores
+        uniform = mx.concatenate([
+            mx.random.uniform(shape=(canvas, global_k),
+                              key=_seed(row.request.seed,
+                                        0xC09A + row.denoise_steps * 1000003))
+            for row in rows
+        ], axis=0)
+        uniform = mx.clip(uniform, 1.17549435e-38, 1.0 - 1.19209290e-7)
+        gumbel = eligible_scores - mx.log(-mx.log(uniform))
+        chosen_local = mx.stop_gradient(mx.argmax(gumbel, axis=-1))
+        chosen_score = mx.take_along_axis(cand_scores, chosen_local[:, None], axis=-1)[:, 0]
+        chosen = mx.take_along_axis(cand_tokens, chosen_local[:, None], axis=-1)[:, 0]
+
+    confidence = mx.clip(mx.exp(chosen_score - log_z), 0.0, 1.0)
+    greedy_confidence = mx.clip(mx.exp(greedy_score - log_z), 0.0, 1.0)
+    entropy = log_z - moment_weighted / mx.maximum(moment_sum, 1.17549435e-38)
+    shape = (batch, canvas)
+    return (chosen.reshape(shape), confidence.reshape(shape), entropy.reshape(shape),
+            greedy.reshape(shape), greedy_confidence.reshape(shape))
+
+
+class MLXContinuousEngine:
+    def __init__(self, runtime: MLXRuntime, *, prefix_mib: int,
+                 prefill_chunk: int, vocab_chunk: int, max_batch_rows: int,
+                 max_batch_tokens: int, emit: Any, pipeline_depth: int = 2,
+                 token_events: bool = True):
+        self.runtime = runtime
+        if min(prefill_chunk, vocab_chunk, max_batch_rows, pipeline_depth) <= 0:
+            raise ValueError("Chunk sizes, batch rows and pipeline depth must be positive.")
+        if max_batch_tokens < int(runtime.config.canvas_length):
+            raise ValueError("--max-batch-tokens must fit at least one canvas.")
+        self.prefix = PrefixKVCache(prefix_mib * 1024 * 1024)
+        self.prefill_chunk = prefill_chunk
+        self.vocab_chunk = vocab_chunk
+        self.max_batch_rows = max_batch_rows
+        self.max_batch_tokens = max_batch_tokens
+        self.emit = emit
+        self.token_events = token_events
+        self.pipeline_depth = pipeline_depth
+        self.forward_count = 0
+        self.active_row_steps = 0
+        self.denoise_seconds = 0.0
+        self.max_inflight = 0
+        self.peak_active_requests = 0
+        generation = runtime.generation
+        pad_token_id = generation.pad_token_id
+        if pad_token_id is None:
+            pad_token_id = getattr(runtime.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)
+        configured_eos = generation.eos_token_id or runtime.config.eos_token_id
+        if isinstance(configured_eos, int):
+            configured_eos = [configured_eos]
+        self.turn_end = int(runtime.config.turn_end_token_id if
+                            generation.turn_end_token_id is None else
+                            generation.turn_end_token_id)
+        self.stops = tuple(dict.fromkeys((self.turn_end, *(int(x) for x in configured_eos or ()))))
+        token_values = [generation.pad_token_id, generation.bos_token_id,
+                        generation.eos_token_id, generation.turn_end_token_id,
+                        getattr(runtime.config, "image_token_id", None),
+                        generation.repetition_penalty_exclude_token_ids]
+        self.excluded = set()
+        for value in token_values:
+            if isinstance(value, int):
+                self.excluded.add(int(value))
+            elif isinstance(value, (tuple, list, set)):
+                self.excluded.update(int(x) for x in value if x is not None)
+
+    def begin_prefill(self, request: ContinuousRequest, created: float) -> PrefillRow:
+        encoded = apply_chat_template(self.runtime.tokenizer, request.messages,
+                                      think=request.think)
+        prompt = tuple(_extract_input_ids(encoded))
+        if len(prompt) + request.max_new_tokens > int(self.runtime.config.text_config.max_position_embeddings):
+            raise ValueError("Prompt plus response exceeds the model position limit.")
+        prefix_length, cache = self.prefix.longest(prompt)
+        if cache is None:
+            cache = self.runtime.model.model.encoder.make_cache()
+        return PrefillRow(request, prompt, cache, prefix_length, prefix_length, created)
+
+    def prefill_step(self, work: PrefillRow) -> InferenceRow | None:
+        """Execute at most one prefill chunk before yielding to active rows."""
+        started = time.perf_counter()
+        if work.offset < len(work.prompt_ids):
+            end = min(work.offset + self.prefill_chunk, len(work.prompt_ids))
+            block = mx.array(work.prompt_ids[work.offset:end], mx.int32)[None, :]
+            _, work.cache = self.runtime.model.model.encoder(block, cache=work.cache)
+            mx.eval(*(value for layer in work.cache for value in (layer.keys, layer.values)
+                      if isinstance(value, mx.array)))
+            work.offset = end
+            self.prefix.put(work.prompt_ids[:end], work.cache)
+        if work.offset < len(work.prompt_ids):
+            work.seconds += time.perf_counter() - started
+            return None
+        request, prompt = work.request, work.prompt_ids
+        cache = _clone_cache(work.cache, compact=True)
+        work.seconds += time.perf_counter() - started
+        seen = {token for token in prompt if token not in self.excluded}
+        row = InferenceRow(request, prompt, cache,
+                           _empty_rolling(self.runtime.config, request.seed,
+                                          request.max_new_tokens, self.pad_token_id), [],
+                           seen, work.created, time.perf_counter(),
+                           prefill_seconds=work.seconds)
+        self.emit({"event": "request_started", "request_id": request.request_id,
+                   "prompt_tokens": len(prompt), "prefix_cache_tokens": work.reused,
+                   "prefill_seconds": row.prefill_seconds})
+        return row
+
+    def admit(self, request: ContinuousRequest, created: float) -> InferenceRow:
+        work = self.begin_prefill(request, created)
+        while True:
+            row = self.prefill_step(work)
+            if row is not None:
+                return row
+
+    def _append_encoder(self, rows: list[InferenceRow], blocks: list[list[int]]) -> None:
+        # Keep the encoder's GEMM and expert routing shapes independent of the
+        # cohort too. Only the scheduler and GPU submission are concurrent.
+        for row, tokens in zip(rows, blocks, strict=True):
+            if tokens:
+                ids = mx.array(tokens, mx.int32)[None, :]
+                _, row.cache = self.runtime.model.model.encoder(ids, cache=row.cache)
+                mx.async_eval(*(value for layer in row.cache
+                                for value in (layer.keys, layer.values)
+                                if isinstance(value, mx.array)))
+
+    def step(self, rows: list[InferenceRow]) -> list[InferenceRow]:
+        if not rows:
+            return []
+        if len(rows) > min(self.max_batch_rows, self.max_batch_tokens //
+                           int(self.runtime.config.canvas_length)):
+            raise ValueError("Denoise cohort exceeds the configured row/token budget.")
+        if len({id(row) for row in rows}) != len(rows):
+            raise ValueError("A request may appear only once in a denoise cohort.")
+        started = time.perf_counter()
+        finished = []
+        for start in range(0, len(rows), self.pipeline_depth):
+            cohort = rows[start:start + self.pipeline_depth]
+            outputs = _forward_independent_rows(self.runtime.model, cohort)
+            pending = [self._prepare_step([row], output=output)
+                       for row, output in zip(cohort, outputs, strict=True)]
+            self.max_inflight = max(self.max_inflight, len(pending))
+            for work in pending:
+                finished.extend(self._finish_step(work))
+        self.denoise_seconds += time.perf_counter() - started
+        return finished
+
+    def _prepare_step(self, rows: list[InferenceRow], *,
+                      output: MLXCanvasOutput) -> _DenoiseWork:
+        model, config, generation = (self.runtime.model, self.runtime.config,
+                                     self.runtime.generation)
+        rolling = _concat_rows([row.rolling for row in rows])
+        batch, canvas = rolling.canvas.shape
+        stats = _inference_statistics(
+            output.heavy_hidden, model.model.decoder.embed_tokens.weight,
+            rows, softcap=float(config.text_config.final_logit_softcapping),
+            chunk_size=self.vocab_chunk, penalty=float(generation.repetition_penalty),
+            excluded=self.excluded,
+            top_k=getattr(config, "commit_top_k", 40),
+            min_p=getattr(config, "commit_min_p", 0.05),
+        )
+        proposal, confidence, entropy, greedy, greedy_confidence = stats
+        confidence = confidence.astype(mx.float32)
+        entropy = entropy.astype(mx.float32)
+        changed = (proposal != rolling.canvas).astype(mx.float32)
+        remaining = mx.array([row.request.max_new_tokens - len(row.generated)
+                              for row in rows], mx.int32)
+        physical_positions = (mx.arange(canvas)[None, :] - rolling.head[:, None]) % canvas
+        valid = physical_positions < remaining[:, None]
+        next_latent = replace(
+            output.next_latent_state, confidence=confidence, entropy=entropy,
+            age=rolling.latent.age + 1, token_changed=changed,
+            confidence_delta=confidence - rolling.latent.confidence,
+            entropy_delta=entropy - rolling.latent.entropy,
+        )
+        next_latent = _concat_rows([
+            model.latent_deliberation.observe_state(
+                _slice_row(next_latent, index), output.heavy_hidden[index:index + 1],
+                output.working_state[index:index + 1], valid[index:index + 1],
+                rolling.head[index:index + 1]) for index in range(batch)
+        ])
+        policy = select_commit_lengths(
+            _logical(proposal, rolling.head),
+            _logical(fused_commit_failure_rate(
+                confidence, entropy, entropy_weight=config.commit_entropy_weight,
+                confidence_power=config.commit_confidence_power,
+                top_k=getattr(config, "commit_top_k", None),
+                min_p=getattr(config, "commit_min_p", None),
+                target_confidence=getattr(config, "commit_target_confidence", None),
+                failure_budget=float(config.commit_failure_budget)), rolling.head),
+            _logical(fused_commit_failure_rate(
+                rolling.latent.confidence, rolling.latent.entropy,
+                entropy_weight=config.commit_entropy_weight,
+                confidence_power=config.commit_confidence_power,
+                top_k=getattr(config, "commit_top_k", None),
+                min_p=getattr(config, "commit_min_p", None),
+                target_confidence=getattr(config, "commit_target_confidence", None),
+                failure_budget=float(config.commit_failure_budget)), rolling.head),
+            _logical(greedy, rolling.head),
+            _logical(fused_commit_failure_rate(
+                greedy_confidence, entropy, entropy_weight=config.commit_entropy_weight,
+                confidence_power=config.commit_confidence_power,
+                top_k=getattr(config, "commit_top_k", None),
+                min_p=getattr(config, "commit_min_p", None),
+                target_confidence=getattr(config, "commit_target_confidence", None),
+                failure_budget=float(config.commit_failure_budget)), rolling.head),
+            ponder_steps=rolling.latent.ponder_steps,
+            stagnation_steps=rolling.latent.stagnation_steps,
+            active_rows=mx.ones((batch,), mx.bool_), remaining_lengths=remaining,
+            failure_budget=float(config.commit_failure_budget),
+            stop_token_id=self.stops,
+            stagnation_threshold=int(generation.jump_on_no_progress_after),
+            min_progress=float(generation.min_trajectory_progress),
+            max_ponder_steps=int(generation.max_ponder_steps),
+            valid_mask=mx.arange(canvas)[None, :] < remaining[:, None],
+        )
+        mx.async_eval(policy.commit_lengths, policy.commit_token_ids,
+                      policy.jump_rows, policy.ponder_steps, policy.stagnation_steps,
+                      *_arrays(next_latent), output.heavy_hidden, output.working_state)
+        self.forward_count += 1
+        return _DenoiseWork(rows, rolling, output, proposal, remaining,
+                            physical_positions, next_latent, policy)
+
+    def _finish_step(self, work: _DenoiseWork) -> list[InferenceRow]:
+        rows, rolling, output = work.rows, work.rolling, work.output
+        proposal, remaining = work.proposal, work.remaining
+        physical_positions, next_latent, policy = (
+            work.physical_positions, work.next_latent, work.policy)
+        model, config, generation = (self.runtime.model, self.runtime.config,
+                                     self.runtime.generation)
+        batch, canvas = rolling.canvas.shape
+        lengths = policy.commit_lengths.astype(mx.int32)
+        selected = policy.commit_token_ids
+        mx.eval(lengths, selected, policy.jump_rows)
+        completed = time.perf_counter()
+        for row in rows:
+            row.last_denoise_at = completed
+            if row.first_denoise_at is None:
+                row.first_denoise_at = completed
+                self.emit({"event": "first_denoise", "request_id": row.request.request_id,
+                           "ttfd_seconds": completed - row.created})
+        lengths_host = [int(x) for x in lengths.tolist()]
+        maximum = max(lengths_host)
+        selected_host = selected.tolist()
+        blocks = [[int(x) for x in selected_host[index][:lengths_host[index]]]
+                  for index in range(batch)]
+        # The commit policy has fixed these tokens. Stream them before the
+        # writer and encoder append, which are only needed by the next step.
+        for index, row in enumerate(rows):
+            block = blocks[index]
+            if not block:
+                continue
+            if row.first_token_at is None:
+                row.first_token_at = time.perf_counter()
+            row.generated.extend(block)
+            row.repetition_seen.update(token for token in block
+                                       if token not in self.excluded)
+            if self.token_events:
+                self.emit({"event": "token", "request_id": row.request.request_id,
+                           "token_ids": block, "generated_tokens": len(row.generated),
+                           "text": self.runtime.tokenizer.decode(
+                               row.generated, skip_special_tokens=False),
+                           "denoise_steps": row.denoise_steps + 1})
+        next_canvas = mx.where(
+            (physical_positions < lengths[:, None]) & policy.jump_rows[:, None],
+            _physical(selected, rolling.head), proposal,
+        )
+        next_latent = replace(next_latent, ponder_steps=policy.ponder_steps,
+                              stagnation_steps=policy.stagnation_steps)
+        if maximum:
+            embeddings = model.model.decoder.embed_tokens(selected[:, :maximum])
+            embeddings = embeddings * model.model.decoder.embed_scale
+            memory, _ = model.latent_deliberation.commit_write(
+                memory=next_latent.memory_slots,
+                working_state=output.working_state,
+
+
+                heavy_hidden=output.heavy_hidden,
+                committed_token_embeddings=embeddings,
+                commit_lengths=lengths,
+                prefix_lengths=mx.array([int(row.cache[0].offset) for row in rows], mx.int32),
+                commit_reason=infer_commit_reason(
+                    lengths, jump_rows=policy.jump_rows,
+                    commit_token_ids=selected, terminal_token_ids=self.stops,
+                ),
+                canvas_head=rolling.head, max_commit=maximum,
+            )
+            next_latent = replace(
+                next_latent, memory_slots=memory,
+                gdn2=replace(next_latent.gdn2, persistent=memory),
+            )
+        self._append_encoder(rows, blocks)
+        noise = mx.concatenate([
+            _noise(row.request.seed, row.denoise_steps + 1, canvas,
+                   int(config.text_config.vocab_size)) for row in rows
+        ], axis=0)
+        next_rolling = MLXRollingState(next_canvas, next_latent,
+                                        rolling.head).advance_ring(
+            lengths, noise, entropy_fill_value=math.log(config.text_config.vocab_size),
+        )
+        logical_positions = (mx.arange(canvas)[None, :] - next_rolling.head[:, None]) % canvas
+        newly_exposed = logical_positions >= (canvas - lengths)[:, None]
+        next_rolling = replace(
+            next_rolling,
+            canvas=mx.where(newly_exposed &
+                            (logical_positions >= (remaining - lengths)[:, None]),
+                            self.pad_token_id, next_rolling.canvas),
+        )
+        # The following step depends on these arrays on the same MLX stream.
+        # Submit now and overlap the writer/refill with host stop/output work.
+        mx.async_eval(*_arrays(next_rolling))
+        finished = []
+        jumped = policy.jump_rows.tolist()
+        for index, row in enumerate(rows):
+            row.rolling = _slice_row(next_rolling, index)
+            row.denoise_steps += 1
+            self.active_row_steps += 1
+            row.jumps += int(jumped[index])
+            row.shifts += int(bool(blocks[index]))
+            if self.turn_end in blocks[index]:
+                row.stop_reason = "turn_end"
+            elif any(token in self.stops for token in blocks[index]):
+                row.stop_reason = "eos"
+            elif len(row.generated) >= row.request.max_new_tokens:
+                row.stop_reason = "max_new_tokens"
+            elif (row.request.max_denoising_steps is not None and
+                  row.denoise_steps >= row.request.max_denoising_steps):
+                row.stop_reason = "max_denoising_steps"
+            else:
+                bound = row.request.max_new_tokens * int(generation.max_ponder_steps)
+                if row.denoise_steps >= bound:
+                    row.stop_reason = "episode_watchdog"
+            if row.stop_reason is not None:
+                finished.append(row)
+        return finished
+
+    def result(self, row: InferenceRow) -> dict[str, Any]:
+        now = time.perf_counter()
+        return {"event": "generation_result", "request_id": row.request.request_id,
+                "status": "finished", "checkpoint": str(self.runtime.checkpoint),
+                "step": self.runtime.step, "thinking": "on" if row.request.think else "off",
+                "stop_reason": row.stop_reason, "prompt_tokens": len(row.prompt_ids),
+                "generated_tokens": len(row.generated), "token_ids": row.generated,
+                "text": self.runtime.tokenizer.decode(row.generated, skip_special_tokens=True),
+                "denoise_steps": row.denoise_steps, "jump_count": row.jumps,
+                "state_shift_count": row.shifts,
+                "tokens_per_forward": len(row.generated) / max(row.denoise_steps, 1),
+                "queue_seconds": row.admitted - row.created,
+                "prefill_seconds": row.prefill_seconds,
+                "ttfd_seconds": None if row.first_denoise_at is None else row.first_denoise_at - row.created,
+                "denoise_steps_per_second": row.denoise_steps / max(
+                    (row.last_denoise_at or now) - row.admitted, 1e-9),
+                "ttft_seconds": None if row.first_token_at is None else row.first_token_at - row.created,
+                "elapsed_seconds": now - row.created,
+                "prefix_cache_hits": self.prefix.hits,
+                "prefix_cache_reused_tokens": self.prefix.reused_tokens,
+                "mlx_active_memory_bytes": int(mx.get_active_memory()),
+                "mlx_peak_memory_bytes": int(mx.get_peak_memory())}
+
+
+def iter_scheduler(engine: MLXContinuousEngine, incoming: queue.Queue[Any],
+                   max_queue_size: int, *,
+                   continue_on_error: bool = False) -> Iterator[dict[str, Any]]:
+    """Bounded continuous admission with a prefill token budget per decode turn."""
+    active: deque[InferenceRow] = deque()
+    prefilling: deque[PrefillRow] = deque()
+    pending: deque[tuple[ContinuousRequest, float]] = deque()
+    seen_ids: set[str] = set()
+    input_done = False
+    started = time.perf_counter()
+    capacity = min(engine.max_batch_rows, engine.max_batch_tokens //
+                   int(engine.runtime.config.canvas_length))
+
+    def accept(item: Any, created: float) -> dict[str, Any] | None:
+        nonlocal input_done
+        if item is None:
+            input_done = True
+        elif isinstance(item, dict):
+            return item
+        elif item.request_id in seen_ids:
+            return {"event": "request_error", "request_id": item.request_id,
+                    "error": "Duplicate request_id."}
+        else:
+            seen_ids.add(item.request_id)
+            pending.append((item, created))
+
+    while not input_done or pending or prefilling or active:
+        while not input_done and len(pending) < max_queue_size:
+            try:
+                error = accept(*incoming.get_nowait())
+                if error is not None:
+                    yield error
+            except queue.Empty:
+                break
+        prefill_budget = engine.prefill_chunk
+        had_active = bool(active)
+        while prefill_budget > 0:
+            if pending and len(active) + len(prefilling) < engine.max_batch_rows:
+                request, created = pending.popleft()
+                try:
+                    # A short newcomer gets one chunk promptly; incomplete
+                    # prompts then rotate in the bounded prefill cohort.
+                    prefilling.appendleft(engine.begin_prefill(request, created))
+                except Exception as error:
+                    yield {"event": "request_error", "request_id": request.request_id,
+                           "error": str(error)}
+                    continue
+            if not prefilling:
+                break
+            work = prefilling.popleft()
+            tokens = min(engine.prefill_chunk, len(work.prompt_ids) - work.offset)
+            if tokens > prefill_budget:
+                # Keep original chunk boundaries and singleton GEMM shapes.
+                prefilling.appendleft(work)
+                break
+            prefill_budget -= tokens
+            try:
+                row = engine.prefill_step(work)
+                if row is None:
+                    prefilling.append(work)
+                else:
+                    active.append(row)
+            except Exception as error:
+                yield {"event": "request_error", "request_id": work.request.request_id,
+                       "error": str(error)}
+            if not had_active and active:
+                # Deliver the first request's initial denoise promptly. Once
+                # decoding, use the remaining token budget to fill short rows.
+                break
+        if active:
+            engine.peak_active_requests = max(getattr(engine, "peak_active_requests", 0),
+                                               len(active) + len(prefilling))
+            selected = [active.popleft() for _ in range(min(capacity, len(active)))]
+            try:
+                finished = engine.step(selected)
+                finished_ids = {id(row) for row in finished}
+                for row in selected:
+                    if id(row) in finished_ids:
+                        yield engine.result(row)
+                    else:
+                        active.append(row)
+            except Exception as error:
+                for row in selected:
+                    yield {"event": "generation_result", "request_id": row.request.request_id,
+                           "status": "failed", "error": str(error)}
+                if not continue_on_error:
+                    raise
+        elif not pending and not prefilling and not input_done:
+            # Wake immediately for new input instead of polling every 10 ms.
+            error = accept(*incoming.get())
+            if error is not None:
+                yield error
+    tail_started = time.perf_counter()
+    mx.synchronize()
+    engine.denoise_seconds += time.perf_counter() - tail_started
+    elapsed = time.perf_counter() - started
+    yield {"event": "inference_summary", "elapsed_seconds": elapsed,
+                 "active_row_denoises": engine.active_row_steps,
+                 "heavy_forward_count": engine.forward_count,
+                 "denoise_steps_per_second": engine.active_row_steps / max(elapsed, 1e-9),
+                 "denoise_scheduler_seconds": engine.denoise_seconds,
+                 "scheduler_denoises_per_second": engine.active_row_steps /
+                     max(engine.denoise_seconds, 1e-9),
+                 "pipeline_depth": engine.pipeline_depth,
+                 "max_inflight_denoises": engine.max_inflight,
+                 "peak_active_requests": getattr(engine, "peak_active_requests", 0),
+                 "mlx_peak_memory_bytes": int(mx.get_peak_memory())}
+
+
+def _run_scheduler(engine: MLXContinuousEngine, incoming: queue.Queue[Any],
+                   max_queue_size: int) -> int:
+    had_error = False
+    for record in iter_scheduler(engine, incoming, max_queue_size):
+        had_error |= record["event"] == "request_error" or record.get("status") == "failed"
+        engine.emit(record)
+    return int(had_error)
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Modilify Mk2 native MLX text generation")
+    parser.add_argument("--model", "--checkpoint", dest="checkpoint",
+                        default=(str(Path(__file__).resolve().parent)
+                                 if (Path(__file__).resolve().parent / "export_manifest.json").is_file()
+                                 else "Modilify/Modilify-Mk2-preview-mlx"),
+                        help="Local model directory or Hugging Face Hub model ID.")
+    parser.add_argument("--prompt", default="Why is the sky blue?")
+    parser.add_argument("--requests-jsonl", metavar="PATH|-",
+                        help="Read JSONL requests continuously; '-' reads stdin.")
+    parser.add_argument("--stream", action="store_true",
+                        help="Stream generated text to stdout for a single --prompt request.")
+    parser.add_argument("--think", type=parse_bool, default=True)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--max-new-tokens", type=int, default=8192)
+    parser.add_argument("--max-denoising-steps", type=int)
+    parser.add_argument("--canvas-length", type=int)
+    parser.add_argument("--repetition-penalty", type=float, default=1.0)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--pipeline-depth", type=int, default=2,
+                        help="Bound in-flight singleton denoises; 1 minimizes activation memory.")
+    parser.add_argument("--max-batch-tokens", type=int, default=1024)
+    parser.add_argument("--prefill-chunk-size", type=int, default=256)
+    parser.add_argument("--vocab-chunk-size", type=int, default=4096)
+    parser.add_argument("--prefix-cache-mib", type=int, default=512)
+    parser.add_argument("--max-queue-size", type=int, default=128)
+    parser.add_argument("--commit-failure-budget", type=float, default=None,
+                        help="Override commit failure budget (defaults to checkpoint config).")
+    parser.add_argument("--commit-top-k", type=int, default=None,
+                        help="Override commit top_k bound (defaults to checkpoint config).")
+    parser.add_argument("--commit-min-p", type=float, default=None,
+                        help="Override commit min_p bound (defaults to checkpoint config).")
+    parser.add_argument("--commit-target-confidence", type=float, default=None,
+                        help="Override commit target confidence (defaults to checkpoint config).")
+    return parser
+
+
+def _validate_args(parser: argparse.ArgumentParser, args: Any) -> None:
+    if args.stream and args.requests_jsonl is not None:
+        parser.error("--stream can only be used with a single --prompt request.")
+    positive = ("max_new_tokens", "batch_size", "max_batch_tokens",
+                "prefill_chunk_size", "vocab_chunk_size", "max_queue_size", "pipeline_depth")
+    for name in positive:
+        if getattr(args, name) <= 0:
+            parser.error(f"--{name.replace('_', '-')} must be positive.")
+    if args.max_denoising_steps is not None and args.max_denoising_steps <= 0:
+        parser.error("--max-denoising-steps must be positive.")
+    if args.prefix_cache_mib < 0:
+        parser.error("--prefix-cache-mib must be nonnegative.")
+    if not math.isfinite(args.repetition_penalty) or args.repetition_penalty <= 0:
+        parser.error("--repetition-penalty must be finite and positive.")
+    if args.commit_failure_budget is not None and args.commit_failure_budget <= 0:
+        parser.error("--commit-failure-budget must be positive.")
+    if args.commit_top_k is not None and args.commit_top_k <= 0:
+        parser.error("--commit-top-k must be positive.")
+    if args.commit_min_p is not None and not (0.0 < args.commit_min_p < 1.0):
+        parser.error("--commit-min-p must be in (0, 1).")
+    if args.commit_target_confidence is not None and not (0.0 < args.commit_target_confidence < 1.0):
+        parser.error("--commit-target-confidence must be in (0, 1).")
+
+
+def _producer(path: str, output: queue.Queue[Any], args: Any) -> None:
+    stream = None
+    try:
+        stream = sys.stdin if path == "-" else open(path, encoding="utf-8")
+        for line_number, line in enumerate(stream, 1):
+            if not line.strip():
+                continue
+            try:
+                record = json.loads(line)
+                request = parse_continuous_request(
+                    record, default_max_new_tokens=args.max_new_tokens,
+                    default_max_denoising_steps=args.max_denoising_steps,
+                    default_seed=args.seed, default_think=args.think,
+                )
+                output.put((request, time.perf_counter()))
+            except Exception as error:
+                output.put(({"event": "request_error", "line": line_number,
+                             "error": str(error)}, time.perf_counter()))
+    except Exception as error:
+        output.put(({"event": "request_error", "error": str(error)},
+                    time.perf_counter()))
+    finally:
+        if stream is not None and stream is not sys.stdin:
+            stream.close()
+        output.put((None, time.perf_counter()))
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = _parser()
+    args = parser.parse_args(argv)
+    _validate_args(parser, args)
+    runtime = load_runtime(args.checkpoint, canvas_length=args.canvas_length,
+                           max_new_tokens=args.max_new_tokens,
+                           max_denoising_steps=args.max_denoising_steps,
+                           repetition_penalty=args.repetition_penalty,
+                           commit_failure_budget=args.commit_failure_budget,
+                           commit_top_k=args.commit_top_k,
+                           commit_min_p=args.commit_min_p,
+                           commit_target_confidence=args.commit_target_confidence)
+    if args.max_denoising_steps is None:
+        args.max_denoising_steps = runtime.generation.max_denoising_steps
+
+    streamed_token_ids: dict[str, list[int]] = {}
+    streamed_text: dict[str, str] = {}
+
+    def emit(record: dict[str, Any]) -> None:
+        if args.stream:
+            event = record.get("event")
+            request_id = str(record.get("request_id", "single"))
+            if event == "token":
+                ids = streamed_token_ids.setdefault(request_id, [])
+                ids.extend(int(token_id) for token_id in record["token_ids"])
+                current = runtime.tokenizer.decode(ids, skip_special_tokens=True)
+                previous = streamed_text.get(request_id, "")
+                if current.startswith(previous):
+                    delta = current[len(previous):]
+                else:
+                    # Keep output append-only if a tokenizer revises a prior decode.
+                    common = 0
+                    for old_char, new_char in zip(previous, current):
+                        if old_char != new_char:
+                            break
+                        common += 1
+                    delta = current[common:]
+                if delta:
+                    sys.stdout.write(delta)
+                    sys.stdout.flush()
+                streamed_text[request_id] = current
+            elif event == "generation_result":
+                if record.get("status") == "failed":
+                    sys.stderr.write(f"\nInference failed: {record.get('error', 'unknown error')}\n")
+                    sys.stderr.flush()
+                else:
+                    sys.stdout.write("\n")
+                    sys.stdout.flush()
+            elif event == "request_error":
+                sys.stderr.write(f"\nInference error: {record.get('error', 'unknown error')}\n")
+                sys.stderr.flush()
+            return
+        sys.stdout.write(json.dumps(record, ensure_ascii=False) + "\n")
+        sys.stdout.flush()
+
+    emit({"event": "runtime_loaded", "checkpoint": str(runtime.checkpoint),
+          "step": runtime.step, "restored_trainable_tensors": runtime.tensor_count,
+          "backend": "mlx", "canvas_length": runtime.config.canvas_length,
+          "pipeline_depth": args.pipeline_depth,
+          "execution": "layer_interleaved_singleton",
+          "commit_failure_budget": runtime.config.commit_failure_budget,
+          "commit_target_confidence": getattr(runtime.config, "commit_target_confidence", None),
+          "commit_top_k": getattr(runtime.config, "commit_top_k", None),
+          "commit_min_p": getattr(runtime.config, "commit_min_p", None)})
+    engine = MLXContinuousEngine(
+        runtime, prefix_mib=args.prefix_cache_mib,
+        prefill_chunk=args.prefill_chunk_size,
+        vocab_chunk=args.vocab_chunk_size,
+        max_batch_rows=args.batch_size,
+        max_batch_tokens=args.max_batch_tokens, emit=emit, pipeline_depth=args.pipeline_depth,
+    )
+    incoming: queue.Queue[Any] = queue.Queue(maxsize=max(2, args.max_queue_size))
+    if args.requests_jsonl is None:
+        request = ContinuousRequest("single", [{"role": "user", "content": args.prompt}],
+                                    args.max_new_tokens, args.max_denoising_steps,
+                                    args.seed, args.think, args.prompt)
+        incoming.put((request, time.perf_counter()))
+        incoming.put((None, time.perf_counter()))
+    else:
+        threading.Thread(target=_producer, args=(args.requests_jsonl, incoming, args),
+                         daemon=True).start()
+    return _run_scheduler(engine, incoming, args.max_queue_size)
+
+
+def load_model(
+    model: str = "Modilify/Modilify-Mk2-preview-mlx", *,
+    canvas_length: int | None = None,
+) -> MLXRuntime:
+    """Load the complete local model or its Hugging Face snapshot once."""
+    return load_runtime(model, canvas_length=canvas_length,
+                        max_new_tokens=256, max_denoising_steps=None,
+                        repetition_penalty=1.0)
+
+
+def generate(
+    runtime: MLXRuntime, prompt: str | None = None, *,
+    messages: list[dict[str, Any]] | None = None,
+    max_new_tokens: int = 256, max_denoising_steps: int | None = None,
+    seed: int = 42, think: bool = True,
+) -> dict[str, Any]:
+    """Generate one independent response, returning text, tokens, and metrics."""
+    record = {"request_id": "single", "max_new_tokens": max_new_tokens,
+              "max_denoising_steps": max_denoising_steps, "seed": seed, "think": think}
+    if prompt is not None:
+        record["prompt"] = prompt
+    if messages is not None:
+        record["messages"] = messages
+    request = parse_continuous_request(
+        record, default_max_new_tokens=256,
+        default_max_denoising_steps=runtime.generation.max_denoising_steps,
+        default_seed=42, default_think=True,
+    )
+    engine = MLXContinuousEngine(
+        runtime, prefix_mib=0, prefill_chunk=256, vocab_chunk=4096,
+        max_batch_rows=1, max_batch_tokens=int(runtime.config.canvas_length),
+        emit=lambda event: None, pipeline_depth=1, token_events=False,
+    )
+    row = engine.admit(request, time.perf_counter())
+    while not engine.step([row]):
+        pass
+    return engine.result(row)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

model-00001.safetensors

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

model-00002.safetensors

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

model-00003.safetensors

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

model-00004.safetensors

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

model-00005.safetensors

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

model-00006.safetensors

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

model-00007.safetensors

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

model-00008.safetensors

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

model-00009.safetensors

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

model-00010.safetensors

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

model-00011.safetensors

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

model-00012.safetensors

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

model-00013.safetensors

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

model-00014.safetensors

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

model-00015.safetensors

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

model-00016.safetensors

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

model-00017.safetensors

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

model-00018.safetensors

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

model-00019.safetensors

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

model-00020.safetensors

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

model-00021.safetensors

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

model-00022.safetensors

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

model-00023.safetensors

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

model-00024.safetensors

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

model-00025.safetensors

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

model-00026.safetensors

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

model-00027.safetensors

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

model-00028.safetensors

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

model-00029.safetensors

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

model-00030.safetensors

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

model-00031.safetensors

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

model-00032.safetensors

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

modilify_mk2/__init__.py

@@ -0,0 +1,9 @@
+"""Modilify Mk2 native MLX inference package."""
+
+from .configuration_modilify_mk2 import ModilifyMk2Config, ModilifyMk2TextConfig
+from transformers import AutoConfig
+
+AutoConfig.register(ModilifyMk2Config.model_type, ModilifyMk2Config, exist_ok=True)
+
+__version__ = "1.0.0"
+__all__ = ["ModilifyMk2Config", "ModilifyMk2TextConfig"]

modilify_mk2/chat.py

@@ -0,0 +1,287 @@
+"""Modilify chat-template rendering for MLX inference."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from typing import Any
+
+GEMMA_THOUGHT_CLOSE = "<channel|>"
+
+_THINK_LINE_RE = re.compile(r"^think(?:\r?\n|$)")
+
+_CHANNEL_BLOCK_RE = re.compile(
+    r"<\|channel>thought\n(.*?)\n?<channel\|>\s*(.*)",
+    flags=re.DOTALL,
+)
+
+_LITERAL_THINK_RE = re.compile(
+    r"\s*<think>(.*?)</think>\s*(.*)",
+    flags=re.DOTALL,
+)
+
+
+def _split_thought_and_content(content: Any) -> tuple[str | None, str]:
+    if not isinstance(content, str):
+        return None, ""
+    text = content.strip()
+    if not text:
+        return None, ""
+    if "�" in text:
+        raise ValueError("Assistant target contains a Unicode replacement character.")
+    if _THINK_LINE_RE.match(text):
+        raise ValueError("Assistant target uses ambiguous literal think without channel markers.")
+    channel_match = _CHANNEL_BLOCK_RE.fullmatch(text)
+    literal_match = _LITERAL_THINK_RE.fullmatch(text)
+    has_channel_token = "<|channel>" in text or GEMMA_THOUGHT_CLOSE in text
+    has_literal_token = "<think>" in text or "</think>" in text
+    if has_channel_token and channel_match is None:
+        raise ValueError("Malformed Gemma channel in assistant target.")
+    if has_literal_token and literal_match is None:
+        raise ValueError("Malformed `<think>` block in assistant target.")
+    if channel_match is not None:
+        thought, answer = channel_match.groups()
+    elif literal_match is not None:
+        thought, answer = literal_match.groups()
+    else:
+        return None, text
+    thought = thought.strip()
+    return (thought or None), answer.strip()
+
+
+def _assistant_thought_and_content(message: dict[str, Any]) -> tuple[str | None, str]:
+    thought, content = _split_thought_and_content(message.get("content"))
+    explicit = message.get("reasoning") or message.get("reasoning_content")
+    if isinstance(explicit, str) and explicit.strip():
+        return explicit.strip(), content
+    return thought, content
+
+
+def _deserialize_tool_call_arguments(arguments: Any) -> dict[str, Any] | None:
+    """Convert OpenAI-style JSON argument strings into the mapping Gemma's template requires."""
+    if arguments is None or isinstance(arguments, dict):
+        return arguments
+    if not isinstance(arguments, str):
+        raise ValueError(
+            "chat_template: tool_calls[].function.arguments must be a JSON object "
+            f"(mapping), not a {type(arguments).__name__}."
+        )
+    text = arguments.strip()
+    if not text:
+        return {}
+    try:
+        parsed = json.loads(text)
+    except json.JSONDecodeError as error:
+        raise ValueError(
+            "chat_template: tool_calls[].function.arguments must be a JSON object "
+            "(mapping), not a string. Deserialize arguments before passing to "
+            f"the template: {error}"
+        ) from error
+    if parsed is None or isinstance(parsed, dict):
+        return parsed
+    raise ValueError(
+        "chat_template: tool_calls[].function.arguments must be a JSON object "
+        f"(mapping), not a {type(parsed).__name__}."
+    )
+
+
+def _stable_tool_call_id(tool_call: dict[str, Any], index: int) -> str:
+    """Create a deterministic id for traces that omitted OpenAI call ids."""
+    payload = json.dumps(
+        tool_call,
+        ensure_ascii=False,
+        sort_keys=True,
+        separators=(",", ":"),
+        default=str,
+    ).encode("utf-8")
+    return f"call_modilify_mk2_{index}_{hashlib.sha1(payload).hexdigest()[:16]}"
+
+
+def _normalize_message_tool_calls(message: dict[str, Any]) -> dict[str, Any]:
+    tool_calls = message.get("tool_calls")
+    if not isinstance(tool_calls, list) or not tool_calls:
+        return message
+    updated_calls = list(tool_calls)
+    changed = False
+    for index, tool_call in enumerate(tool_calls):
+        if not isinstance(tool_call, dict):
+            continue
+        function = tool_call.get("function")
+        # Some agent traces use the compact {name, arguments} shape instead
+        # of OpenAI's {function: {name, arguments}} wrapper.
+        if not isinstance(function, dict):
+            name = tool_call.get("name")
+            if not isinstance(name, str) or not name.strip():
+                continue
+            function = {
+                "name": name,
+                "arguments": tool_call.get(
+                    "arguments", tool_call.get("input", {})
+                ),
+            }
+            changed = True
+        arguments = function.get("arguments")
+        parsed = (
+            arguments
+            if arguments is None or isinstance(arguments, dict)
+            else _deserialize_tool_call_arguments(arguments)
+        )
+        new_function = dict(function)
+        if parsed is not arguments:
+            changed = True
+        new_function["arguments"] = parsed
+        new_call = dict(tool_call)
+        if not isinstance(new_call.get("id"), str) or not new_call["id"]:
+            new_call["id"] = _stable_tool_call_id(tool_call, index)
+            changed = True
+        new_call.setdefault("type", "function")
+        new_call["function"] = new_function
+        updated_calls[index] = new_call
+    if not changed:
+        return message
+    updated = dict(message)
+    updated["tool_calls"] = updated_calls
+    return updated
+
+
+def _normalize_assistant_message(message: dict[str, Any]) -> dict[str, Any]:
+    """Lift think/channel text into ``reasoning`` and deserialize tool arguments."""
+    updated = _normalize_message_tool_calls(message)
+    if updated.get("role") != "assistant":
+        return updated
+    thought, content = _assistant_thought_and_content(updated)
+    content_changed = content != (updated.get("content") or "")
+    reasoning = updated.get("reasoning")
+    needs_reasoning = bool(thought) and reasoning != thought
+    if not content_changed and not needs_reasoning:
+        return updated
+    if updated is message:
+        updated = dict(message)
+    else:
+        updated = dict(updated)
+    if thought:
+        updated["reasoning"] = thought
+    updated["content"] = content
+    return updated
+
+
+def normalize_chat_template_messages(messages: Any) -> Any:
+    """Copy conversations into the official Gemma chat-template message schema."""
+    if not isinstance(messages, list) or not messages:
+        return messages
+    if isinstance(messages[0], list):
+        normalized_batch = None
+        for index, conversation in enumerate(messages):
+            normalized = normalize_chat_template_messages(conversation)
+            if normalized is conversation:
+                continue
+            if normalized_batch is None:
+                normalized_batch = list(messages)
+            normalized_batch[index] = normalized
+        return messages if normalized_batch is None else normalized_batch
+    normalized_messages = None
+    for index, message in enumerate(messages):
+        if not isinstance(message, dict):
+            continue
+        updated = _normalize_assistant_message(message)
+        if updated is message:
+            continue
+        if normalized_messages is None:
+            normalized_messages = list(messages)
+        normalized_messages[index] = updated
+    return messages if normalized_messages is None else normalized_messages
+
+
+def normalize_tool_definitions(tools: Any) -> list[dict[str, Any]] | None:
+    """Normalize optional tool declarations and ignore trace-only tool metadata.
+
+    The native template accepts OpenAI declarations only.  ``data-new`` also
+    contains JSON-encoded declarations and trace metadata shaped like
+    ``{name, arguments, tool_call_id}``; the latter are executed calls, not
+    declarations, and must not be passed to ``format_function_declaration``.
+    """
+    if tools is None:
+        return None
+    pending: list[Any]
+    if isinstance(tools, str):
+        try:
+            parsed = json.loads(tools)
+        except json.JSONDecodeError:
+            return None
+        pending = parsed if isinstance(parsed, list) else [parsed]
+    elif isinstance(tools, dict):
+        pending = [tools]
+    elif isinstance(tools, list):
+        pending = list(tools)
+    else:
+        return None
+
+    normalized: list[dict[str, Any]] = []
+    for item in pending:
+        if isinstance(item, str):
+            try:
+                item = json.loads(item)
+            except json.JSONDecodeError:
+                continue
+            if isinstance(item, list):
+                pending.extend(item)
+                continue
+        if not isinstance(item, dict):
+            continue
+        function = item.get("function")
+        if isinstance(function, dict):
+            name = function.get("name")
+            if not isinstance(name, str) or not name.strip():
+                continue
+            declaration = dict(function)
+            declaration["description"] = declaration.get("description", "")
+            declaration["parameters"] = declaration.get("parameters") or {}
+            normalized.append({
+                "type": "function",
+                "function": declaration,
+            })
+            continue
+        # Accept the common Anthropic/tool-schema spelling when it really is
+        # a declaration.  Execution records with only `arguments` are skipped.
+        name = item.get("name")
+        parameters = item.get("parameters", item.get("input_schema"))
+        if isinstance(name, str) and name.strip() and isinstance(parameters, dict):
+            normalized.append({
+                "type": "function",
+                "function": {
+                    "name": name,
+                    "description": item.get("description", ""),
+                    "parameters": parameters,
+                },
+            })
+    return normalized or None
+
+
+def apply_chat_template(
+    processor: Any,
+    messages: Any,
+    *,
+    think: bool,
+    return_tensors: str | None = None,
+    padding: bool | str = False,
+    tools: Any = None,
+) -> Any:
+    """Render conversations with the Modilify tokenizer template."""
+
+    template_kwargs: dict[str, Any] = {
+        "tokenize": True,
+        "add_generation_prompt": True,
+        "enable_thinking": think,
+        "return_dict": True,
+    }
+    if return_tensors is not None:
+        template_kwargs["return_tensors"] = return_tensors
+    if padding:
+        template_kwargs["padding"] = padding
+    tools = normalize_tool_definitions(tools)
+    if tools:
+        template_kwargs["tools"] = tools
+    messages = normalize_chat_template_messages(messages)
+    encoded = processor.apply_chat_template(messages, **template_kwargs)
+    return encoded

modilify_mk2/configuration_modilify_mk2.py

@@ -0,0 +1,442 @@
+"""Strict text-only configuration for the schema25 GDN2 protocol."""
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+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}."
+            )
+        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",
+]
+
+
+class ModilifyMk2GenerationConfig:
+    """Native settings; no autoregressive sampler or Torch runtime is needed."""
+
+    def __init__(self, **kwargs: Any):
+        unsupported = (
+            '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',
+        )
+        configured = [name for name in unsupported if kwargs.pop(name, None) not in (None, False)]
+        if configured:
+            raise ValueError(f'Unsupported generation fields: {configured}')
+        defaults = {
+            'max_new_tokens': None, 'max_denoising_steps': None,
+            'bos_token_id': None, 'pad_token_id': None,
+            'eos_token_id': None, 'turn_end_token_id': None, 'max_ponder_steps': 64,
+            'jump_on_no_progress_after': 12, 'min_trajectory_progress': 0.005,
+            'repetition_penalty': 1.0, 'repetition_penalty_exclude_token_ids': [],
+        }
+        for name, default in defaults.items():
+            setattr(self, name, kwargs.pop(name, default))
+        self.repetition_penalty_exclude_token_ids = list(dict.fromkeys(
+            int(value) for value in self.repetition_penalty_exclude_token_ids or ()
+        ))
+
+
+    def validate(self) -> None:
+        for name in ('max_new_tokens', 'max_denoising_steps',
+                     'max_ponder_steps', 'jump_on_no_progress_after'):
+            value = getattr(self, name)
+            if value is not None and (isinstance(value, bool) or not isinstance(value, int) or value <= 0):
+                raise ValueError(f'{name} must be a positive integer.')
+        if not math.isfinite(self.repetition_penalty) or self.repetition_penalty <= 0:
+            raise ValueError('repetition_penalty must be finite and positive.')
+        if not math.isfinite(self.min_trajectory_progress) or self.min_trajectory_progress < 0:
+            raise ValueError('min_trajectory_progress must be finite and nonnegative.')
+        if self.turn_end_token_id is not None and self.turn_end_token_id < 0:
+            raise ValueError('turn_end_token_id must be nonnegative.')
+        if any(isinstance(value, bool) or not isinstance(value, int) or value < 0
+               for value in self.repetition_penalty_exclude_token_ids):
+            raise ValueError('Excluded token IDs must be nonnegative integers.')
+
+
+def load_generation_config(model_dir: str | Path) -> ModilifyMk2GenerationConfig:
+    path = Path(model_dir) / 'generation_config.json'
+    if not path.is_file():
+        raise FileNotFoundError(f'Missing generation configuration: {path}')
+    return ModilifyMk2GenerationConfig(**json.loads(path.read_text(encoding='utf-8')))
+
+
+def configure_generation_config(
+    generation_config: ModilifyMk2GenerationConfig, processor: Any, *,
+    max_new_tokens: int, max_denoising_steps: int | None,
+    repetition_penalty: float | None = None,
+) -> ModilifyMk2GenerationConfig:
+    if generation_config.max_denoising_steps is None:
+        generation_config.max_denoising_steps = 48
+    generation_config.max_new_tokens = max_new_tokens
+    if max_denoising_steps is not None:
+        generation_config.max_denoising_steps = max_denoising_steps
+    if repetition_penalty is not None:
+        generation_config.repetition_penalty = repetition_penalty
+    tokenizer = getattr(processor, 'tokenizer', processor)
+    if generation_config.bos_token_id is None:
+        generation_config.bos_token_id = tokenizer.bos_token_id
+    if generation_config.pad_token_id is None:
+        generation_config.pad_token_id = tokenizer.pad_token_id
+    turn_id = tokenizer.convert_tokens_to_ids('<turn|>')
+    if turn_id is None or turn_id == getattr(tokenizer, 'unk_token_id', None):
+        raise ValueError('Tokenizer must define <turn|>.')
+    generation_config.turn_end_token_id = int(turn_id)
+    if generation_config.eos_token_id is None:
+        generation_config.eos_token_id = [value for value in [tokenizer.eos_token_id] if value is not None]
+    tool_id = tokenizer.convert_tokens_to_ids('<|tool_response>')
+    if tool_id is not None and tool_id != getattr(tokenizer, 'unk_token_id', None):
+        eos = generation_config.eos_token_id
+        eos = [eos] if isinstance(eos, int) else list(eos or ())
+        if int(tool_id) not in eos:
+            generation_config.eos_token_id = [*eos, int(tool_id)]
+    generation_config.repetition_penalty_exclude_token_ids = list(dict.fromkeys(
+        int(value) for value in [*generation_config.repetition_penalty_exclude_token_ids,
+                                *(getattr(tokenizer, 'all_special_ids', None) or ())]
+        if value is not None
+    ))
+    generation_config.validate()
+    return generation_config

modilify_mk2/mlx_commit_policy.py

@@ -0,0 +1,255 @@
+"""Native MLX schema25 confidence fusion and sample-independent commit policy."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import mlx.core as mx
+
+
+JUMP_FAILURE_BUDGET = 2.0
+FUSED_EPS = 1.0e-6
+COMMIT_REASON_NONE = 0
+COMMIT_REASON_NORMAL = 1
+COMMIT_REASON_FORCED_JUMP = 2
+COMMIT_REASON_TERMINAL = 3
+
+
+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: mx.array,
+    token_entropy: mx.array,
+    *,
+    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,
+) -> mx.array:
+    """Match the schema25 excess-entropy sigmoid and checkpoint power."""
+
+    p = mx.clip(mx.nan_to_num(proposal_confidence.astype(mx.float32), nan=0.5), eps, 1.0 - eps)
+    entropy = mx.maximum(mx.nan_to_num(token_entropy.astype(mx.float32), nan=0.0), 0.0)
+    binary_entropy = -p * mx.log(p) - (1.0 - p) * mx.log1p(-p)
+    excess = mx.maximum(entropy - binary_entropy, 0.0)
+    k_eff = None
+    if top_k is not None and top_k > 0:
+        k_eff = mx.full(p.shape, float(top_k), dtype=mx.float32)
+    if min_p is not None and min_p > 0:
+        thresh = mx.maximum(float(min_p) * p, 1.0e-6)
+        k_min_p = mx.maximum((1.0 - p) / thresh, 1.0)
+        k_eff = k_min_p if k_eff is None else mx.minimum(k_eff, k_min_p)
+    if k_eff is not None:
+        max_excess = (1.0 - p) * mx.log(k_eff)
+        excess = mx.minimum(excess, max_excess)
+    bias = commit_target_confidence_bias(target_confidence, failure_budget=failure_budget)
+    fused_logit = mx.log(p) - mx.log1p(-p) + bias - entropy_weight * excess
+    return mx.clip(mx.sigmoid(fused_logit) ** confidence_power, eps, 1.0 - eps)
+
+
+def fused_commit_failure_rate(
+    proposal_confidence: mx.array, token_entropy: mx.array, **kwargs: object
+) -> mx.array:
+    return 1.0 - fused_commit_confidence(proposal_confidence, token_entropy, **kwargs)
+
+
+def _terminal_mask(tokens: mx.array, stop_token_id: int | Sequence[int]) -> mx.array:
+    ids = (stop_token_id,) if isinstance(stop_token_id, int) else tuple(dict.fromkeys(stop_token_id))
+    if not ids:
+        raise ValueError("At least one stop token ID is required.")
+    matches = tokens == int(ids[0])
+    for value in ids[1:]:
+        matches = matches | (tokens == int(value))
+    return matches
+
+
+def prefix_failure_commit_lengths(
+    failure_rate: mx.array,
+    *,
+    failure_budget: float,
+    valid_mask: mx.array | None = None,
+) -> mx.array:
+    """Longest contiguous valid prefix whose cumulative failure stays below 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 = mx.ones(failure_rate.shape, mx.bool_)
+    if valid_mask.shape != failure_rate.shape:
+        raise ValueError("Commit validity mask must match failure rate.")
+    risk = mx.clip(failure_rate.astype(mx.float32), 0.0, 1.0) * valid_mask.astype(mx.float32)
+    allowed = (mx.cumsum(risk, axis=-1) < failure_budget) & (
+        mx.cumprod(valid_mask.astype(mx.int32), axis=-1).astype(mx.bool_)
+    )
+    return mx.sum(mx.cumprod(allowed.astype(mx.int32), axis=-1), axis=-1)
+
+
+def first_committed_token_lengths(
+    proposal: mx.array,
+    commit_lengths: mx.array,
+    token_id: int | Sequence[int],
+    *,
+    positions: mx.array | None = None,
+) -> mx.array:
+    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 = mx.arange(proposal.shape[1])[None, :]
+    elif positions.shape != (1, proposal.shape[1]):
+        raise ValueError("Commit positions must have shape [1, canvas].")
+    matches = _terminal_mask(proposal, token_id) & (positions < commit_lengths[:, None])
+    first = mx.min(mx.where(matches, positions, proposal.shape[1]), axis=-1)
+    return mx.minimum(mx.where(first < proposal.shape[1], first + 1, commit_lengths), commit_lengths)
+
+
+def bounded_prefix_failure_commit_lengths(
+    committed_token_ids: mx.array,
+    failure_rate: mx.array,
+    *,
+    failure_budget: float,
+    remaining_lengths: mx.array,
+    stop_token_id: int | Sequence[int],
+    valid_mask: mx.array | None = None,
+    positions: mx.array | None = None,
+) -> mx.array:
+    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].")
+    lengths = prefix_failure_commit_lengths(
+        failure_rate, failure_budget=failure_budget, valid_mask=valid_mask
+    )
+    lengths = mx.minimum(lengths, mx.maximum(remaining_lengths, 0))
+    return first_committed_token_lengths(
+        committed_token_ids, lengths, stop_token_id, positions=positions
+    )
+
+
+@dataclass(frozen=True)
+class MLXCommitPolicyDecision:
+    normal_lengths: mx.array
+    commit_lengths: mx.array
+    commit_token_ids: mx.array
+    jump_rows: mx.array
+    ponder_steps: mx.array
+    stagnation_steps: mx.array
+
+
+def select_commit_lengths(
+    sampled_token_ids: mx.array,
+    normal_failure_rate: mx.array,
+    previous_failure_rate: mx.array,
+    greedy_token_ids: mx.array,
+    jump_failure_rate: mx.array,
+    *,
+    ponder_steps: mx.array,
+    stagnation_steps: mx.array,
+    active_rows: mx.array,
+    remaining_lengths: mx.array,
+    failure_budget: float,
+    stop_token_id: int | Sequence[int],
+    stagnation_threshold: int,
+    min_progress: float,
+    max_ponder_steps: int | None = None,
+    valid_mask: mx.array | None = None,
+) -> MLXCommitPolicyDecision:
+    """Select normal commits or bounded greedy JUMP independently per row."""
+
+    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].")
+    if not (ponder_steps.shape == stagnation_steps.shape == active_rows.shape
+            == remaining_lengths.shape == sampled_token_ids.shape[:1]):
+        raise ValueError("Commit row inputs must share [batch].")
+    if min_progress < 0:
+        raise ValueError("Minimum progress must be nonnegative.")
+    canvas = normal_failure_rate.shape[1]
+    positions = mx.arange(canvas)[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_failure_commit_lengths(
+        previous_failure_rate, failure_budget=failure_budget, valid_mask=valid_mask
+    )
+    frontier = mx.maximum(previous, normal) + 1
+    valid_lengths = (mx.sum(valid_mask.astype(mx.int32), axis=-1) if valid_mask is not None
+                     else mx.full(frontier.shape, canvas, mx.int32))
+    frontier = mx.minimum(frontier, valid_lengths)
+    progress_mask = (positions < frontier[:, None]) & active_rows[:, None]
+    if valid_mask is not None:
+        progress_mask = progress_mask & valid_mask
+    weights = progress_mask.astype(mx.float32)
+    progress = mx.sum((previous_failure_rate.astype(mx.float32)
+                       - normal_failure_rate.astype(mx.float32)) * weights, axis=-1) / mx.maximum(
+                           mx.sum(weights, axis=-1), 1.0)
+    waiting = active_rows & (normal == 0)
+    next_ponder = mx.where(normal > 0, 0, ponder_steps + waiting.astype(mx.int32))
+    next_stagnation = mx.where(normal > 0, 0, stagnation_steps + waiting.astype(mx.int32))
+    jump = (normal == 0) & active_rows & (next_stagnation >= stagnation_threshold)
+    jump = jump & (progress <= min_progress)
+    if max_ponder_steps is not None and max_ponder_steps > 0:
+        jump = jump | ((normal == 0) & active_rows & (next_ponder >= 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,
+    )
+    commit_token_ids = mx.where(jump[:, None], greedy_token_ids, sampled_token_ids)
+    committed = mx.where(active_rows, mx.where(jump, jump_commit, normal), 0)
+    jump = jump & (committed > 0)
+    next_ponder = mx.where(committed > 0, 0, next_ponder).astype(mx.int32)
+    next_stagnation = mx.where(committed > 0, 0, next_stagnation).astype(mx.int32)
+    return MLXCommitPolicyDecision(
+        normal_lengths=normal,
+        commit_lengths=committed,
+        commit_token_ids=commit_token_ids,
+        jump_rows=jump,
+        ponder_steps=next_ponder,
+        stagnation_steps=next_stagnation,
+    )
+
+
+def infer_commit_reason(
+    commit_lengths: mx.array,
+    *,
+    jump_rows: mx.array | None = None,
+    commit_token_ids: mx.array | None = None,
+    terminal_token_ids: Sequence[int] = (),
+) -> mx.array:
+    """Reason codes consumed by the persistent writer at commit only."""
+
+    committed = commit_lengths > 0
+    reasons = mx.where(committed, COMMIT_REASON_NORMAL, COMMIT_REASON_NONE)
+    if jump_rows is not None:
+        reasons = mx.where(committed & jump_rows, COMMIT_REASON_FORCED_JUMP, reasons)
+    if commit_token_ids is not None and terminal_token_ids:
+        positions = mx.arange(commit_token_ids.shape[1])[None, :]
+        terminal = mx.any(
+            _terminal_mask(commit_token_ids, terminal_token_ids)
+            & (positions < commit_lengths[:, None]), axis=-1,
+        )
+        reasons = mx.where(committed & terminal, COMMIT_REASON_TERMINAL, reasons)
+    return reasons.astype(mx.int32)

modilify_mk2/mlx_gdn2_memory.py

@@ -0,0 +1,117 @@
+"""MLX counterpart of the FP32 GDN2 trajectory recurrence."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+from mlx import nn
+
+
+def _l2(value: mx.array) -> mx.array:
+    return value * mx.rsqrt(mx.maximum(mx.sum(mx.square(value), axis=-1, keepdims=True), 1e-12))
+
+
+class _GateProjection(nn.Module):
+    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 __call__(self, source: mx.array) -> mx.array:
+        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 = mx.zeros((heads,), mx.float32)
+        self.dt_bias = mx.full((heads, key_dim), -6.906255, mx.float32)
+
+    def empty(self, *leading: int) -> mx.array:
+        return mx.zeros((*leading, self.heads, self.key_dim, self.value_dim), mx.float32)
+
+    def _projections(self, source: mx.array):
+        normalized = source.astype(mx.float32)
+        source = (normalized * mx.rsqrt(mx.mean(mx.square(normalized), axis=-1, keepdims=True) + 1e-6)).astype(source.dtype)
+        shape = source.shape[:-1]
+        key_shape = (*shape, self.heads, self.key_dim)
+        value_shape = (*shape, self.heads, self.value_dim)
+        k = _l2(nn.silu(self.k_proj(source).astype(mx.float32)).reshape(key_shape))
+        v = nn.silu(self.v_proj(source).astype(mx.float32)).reshape(value_shape)
+        head_rate = mx.exp(self.a_log.astype(mx.float32)).reshape(
+            *((1,) * (source.ndim - 1)), self.heads, 1)
+        decay = mx.exp(-head_rate * nn.softplus(
+            self.f_proj(source).astype(mx.float32).reshape(key_shape) + self.dt_bias.astype(mx.float32)))
+        erase = mx.sigmoid(self.b_proj(source).astype(mx.float32).reshape(key_shape))
+        write = mx.sigmoid(self.w_proj(source).astype(mx.float32).reshape(value_shape))
+        return k, v, decay, erase, write
+
+    def _query(self, source: mx.array) -> mx.array:
+        return _l2(nn.silu(self.q_proj(source).astype(mx.float32)).reshape(
+            *source.shape[:-1], self.heads, self.key_dim))
+
+    def _output(self, value: mx.array, source: mx.array) -> mx.array:
+        gate = nn.silu(self.g_proj(source).astype(mx.float32).reshape(value.shape))
+        value = value * mx.rsqrt(mx.mean(mx.square(value), axis=-1, keepdims=True) + 1e-6) * gate
+        return self.o_proj(value.reshape(*source.shape[:-1], -1).astype(source.dtype))
+
+    def read(self, state: mx.array, source: mx.array) -> mx.array:
+        if state.shape != (*source.shape[:-1], self.heads, self.key_dim, self.value_dim):
+            raise ValueError("GDN2 state and query leading dimensions differ.")
+        value = mx.sum(self._query(source)[..., None] * state.astype(mx.float32), axis=-2)
+        return self._output(value, source)
+
+    def read_shared(self, state: mx.array, source: mx.array) -> mx.array:
+        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(0, 2, 1, 3)
+        value = (query @ state.astype(mx.float32)).transpose(0, 2, 1, 3)
+        return self._output(value, source)
+
+    @staticmethod
+    def _transition(state, k, v, decay, erase, write, valid):
+        decayed = state.astype(mx.float32) * decay[..., None]
+        old = mx.sum((erase * k)[..., None] * decayed, axis=-2)
+        candidate = decayed + k[..., None] * (write * v - old)[..., None, :]
+        return candidate if valid is None else mx.where(
+            valid[..., None, None, None], candidate, state.astype(mx.float32))
+
+    def transition(self, state: mx.array, source: mx.array,
+                   valid: mx.array | None = None) -> mx.array:
+        # Do not override nn.Module.update: it installs parameters for optimizers,
+        # dtype conversion, and autodiff. State evolution is a separate operation.
+        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: mx.array, source: mx.array,
+                       valid: mx.array) -> mx.array:
+        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"]

modilify_mk2/mlx_gdn2_trajectory.py

@@ -0,0 +1,87 @@
+"""MLX dual-timescale matrix state with denoise and commit lifetimes."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import mlx.core as mx
+from mlx import nn
+
+from .mlx_gdn2_memory import GDN2Memory
+
+
+@dataclass(frozen=True)
+class GDN2TrajectoryState:
+    cells: mx.array
+    row: mx.array
+    persistent: mx.array
+    seen: mx.array
+
+
+    def clear_refill(self, lengths: mx.array, head: mx.array) -> "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 = mx.arange(canvas)[None, :]
+        overwritten = ((physical - head[:, None]) % canvas) < lengths[:, None]
+        return GDN2TrajectoryState(
+            mx.where(overwritten[..., None, None, None], 0.0, self.cells),
+            self.row, self.persistent, self.seen & ~overwritten,
+        )
+
+
+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)
+        self.probe_embed.weight = mx.random.normal((probes, width)) * 0.02
+
+    def empty(self, batch: int, canvas: int) -> GDN2TrajectoryState:
+        return GDN2TrajectoryState(
+            self.cell.empty(batch, canvas), self.row.empty(batch),
+            self.persistent.empty(batch), mx.zeros((batch, canvas), mx.bool_),
+        )
+
+    def read(self, state: GDN2TrajectoryState, query: mx.array) -> mx.array:
+        result = (self.cell.read(state.cells, query)
+                  + self.row.read_shared(state.row, query)
+                  + self.persistent.read_shared(state.persistent, query))
+        return mx.where(state.seen[..., None], result, 0.0)
+
+    def observe(self, state: GDN2TrajectoryState, observation: mx.array,
+                live: mx.array, head: mx.array) -> 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] + mx.arange(canvas)[None, :]) % canvas
+        logical = mx.take_along_axis(observation,
+                                     mx.broadcast_to(logical_idx[..., None], observation.shape), axis=1)
+        logical_live = mx.take_along_axis(live, mx.stop_gradient(logical_idx), axis=1)
+        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 = mx.sum(selected.astype(mx.float32), axis=1, keepdims=True)
+            pooled = mx.sum(logical[:, lo:hi].astype(mx.float32)
+                            * selected[..., None], axis=1)
+            pooled = (pooled / mx.maximum(count, 1.0)).astype(observation.dtype)
+            pooled = pooled + self.probe_embed.weight[probe].astype(observation.dtype)
+            row_state = self.row.transition(row_state, pooled, count[:, 0] > 0)
+        return GDN2TrajectoryState(cells, row_state, state.persistent, seen)
+
+
+__all__ = ["GDN2TrajectoryMemory", "GDN2TrajectoryState"]