Add model card, configuration, and oMLX support

.gitattributes

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

LICENSE

@@ -0,0 +1,16 @@
+Qwen Community License 1.0
+
+Copyright (c) 2026 Qwen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software, including the model weights, parameters, configuration files, inference code and associated documentation files (collectively, the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, sell, deploy, host, fine-tune, and create derivative works from (collectively, "Use" or "Using") copies of the Software; and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+1. The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. If the Software (or any derivative works thereof) is Used for any of the licensee's commercial products or services that have more than 100,000,000 monthly active users or US$ 20,000,000 (or equivalent in other currencies) monthly revenue, respective model name must be prominently displayed on the user interface of such product or service; and,
+
+2. If the licensee or any of its affiliates conducts a Model as a Service or AI Work Assistant business, the licensee shall obtain a separate license from Qwen before Using the Software or its derivative works for any commercial purpose. The foregoing requirement shall not apply to the licensee's internal Use of the Software, provided that such Use does not make the Software, its outputs, or its underlying model capabilities available to any third party.
+
+"Model as a Service" means giving a third party access to language model inference or fine-tuning (e.g., via API or a hosted endpoint) in a manner that allows such third parties to exercise meaningful control over the inputs, parameters, or training data. This does not include the mere relaying of requests to models hosted by other third parties. 
+“AI Work Assistant” means an independent AI-powered product primarily designed for AI-assisted coding or office productivity (e.g., Qoder and QwenWork). It does not include: (a) a single-purpose AI tool (such as an AI translation tool); (b) an AI assistant primarily designed for a domain other than coding or office productivity (such as Taobao AI Shopping Assistant or AMap AI Chat); or (c) an AI assistant that is a feature of a product whose primary purpose is not AI-assisted coding or office productivity.
+
+THE SOFTWARE AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL QWEN, ITS AFFILIATES OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. THE USE OF THE SOFTWARE MUST COMPLY WITH APPLICABLE LAWS AND REGULATIONS, AND MUST NOT INFRINGE THE INTELLECTUAL PROPERTY RIGHTS OF ANY THIRD PARTY.
+
+For any questions regarding this license, please contact model-business@notice.qwencloud.com.

README.md

@@ -0,0 +1,157 @@
+---
+base_model:
+- Qwen/Qwen3.8-Flash-Next-FP8
+base_model_relation: quantized
+library_name: mlx
+license: other
+license_name: qwen-community-1.0
+license_link: LICENSE
+pipeline_tag: text-generation
+tags:
+- mlx
+- safetensors
+- qwen4_exp
+- omlx
+- oq
+- quantized
+- apple-silicon
+- mixture-of-experts
+- 4-bit
+- long-context
+- conversational
+- tool-calling
+- speculative-decoding
+- mtp
+---
+
+# Qwen3.8 Flash Next oQ4e MTP 128K for oMLX
+
+This is a community, text-only dynamic quantization of [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8) for Apple Silicon and oMLX.
+
+It is designed for a 131,072-token context window on a Mac with 128 GiB of unified memory. The conversion removes the vision encoder and retains the released one-layer Multi-Token Prediction head for depth-1 speculative decoding. It is not an official Qwen release.
+
+## Quantization
+
+The model uses oMLX enhanced importance-matrix quantization with 4-bit affine, group-size 128 weights as the default. Precision is assigned by tensor family:
+
+- The token embedding and language-model head use 4-bit affine quantization.
+- Large linear and routed-expert weights use importance-matrix-guided mixed precision.
+- Sensitive attention and shared-expert matrices retain higher precision, including 8-bit floors.
+- Routers, recurrent state, convolutions, normalization, residual controls, and related control tensors retain BF16.
+- The MTP head's 512 routed experts use importance-weighted 4-bit affine quantization at group size 128. Their three projections are stacked, avoiding 3,072 separate expert weight and scale arrays.
+- MTP attention, shared-expert, and layer mixer matrices use 8-bit affine quantization. Fusion projections, the final mixer, routers, norms, gates, and injection controls retain BF16.
+- The 51.2B-parameter PLE n-gram bank remains split into 128 shards. Shards 0 through 3 use 2-bit affine, group size 32. Shards 4 through 127 use 3-bit affine, group size 32.
+
+The importance matrix was collected from 1,024 tool-heavy multilingual calibration samples at sequence length 512. The included `oq_imatrix_report.json` records the allocation inputs and coverage.
+
+The MTP expert quantizer uses the calibrated activation importance of the analogous final backbone layer as a proxy. This preserves expert-specific weighting without claiming that the non-MTP calibration cache directly measured the draft head.
+
+The resulting checkpoint contains 22 safetensors files and occupies 87.959 GiB. Its index contains 2,916 tensors, including 62 MTP tensors.
+
+## Runtime requirements
+
+This checkpoint is not compatible with stock Transformers, stock MLX-LM, or an oMLX release without Qwen4 experimental support. The exact text model, cache integration, and mixed-bit SSD-backed PLE implementation used for validation are included under `omlx_support`.
+
+The tested software stack is:
+
+- Apple Silicon
+- MLX and mlx-metal 0.32.1
+- The MLX-LM package bundled with oMLX 0.6.3rc3
+- oMLX 0.6.3rc3
+- Depth-1 MTP enabled through oMLX's generic native-MTP scheduler
+
+The Qwen4 implementation is based on [oMLX pull request #3161](https://github.com/jundot/omlx/pull/3161) at commit `79e11cf86360ca800db5bdc45cec4ee582ada6b1`, with an adaptation for this checkpoint's per-shard 2-bit and 3-bit PLE metadata. The MTP equations follow draft [pull request #3163](https://github.com/jundot/omlx/pull/3163) at commit `b115fa5d928c11acf4588a55814a307d9161baff`, adapted to the sparse 128K text path and oMLX's installed generic speculative scheduler.
+
+Install the tested MLX runtime inside the downloaded model directory, then launch the bundled oMLX wrapper:
+
+```sh
+hf download jedisct1/Qwen3.8-Flash-Next-oQ4e-MTP-128k \
+  --local-dir models/Qwen3.8-Flash-Next-oQ4e-MTP-128k
+cd models/Qwen3.8-Flash-Next-oQ4e-MTP-128k
+uv pip install --target .mlx-runtime mlx==0.32.1 mlx-metal==0.32.1
+zsh omlx_support/serve
+```
+
+The launcher expects oMLX at `/Applications/oMLX.app`, listens only on `127.0.0.1:8766`, and uses `omlx` as its default local API key. Set `OMLX_APP_RESOURCES`, `OMLX_API_KEY`, `OMLX_PORT`, `OMLX_CACHE_DIR`, or `OMLX_MLX_RUNTIME` to override those paths and values.
+
+It installs the included model settings into a fresh oMLX base path, permits one concurrent request, selects oMLX's balanced memory guard, disables the in-memory hot cache, memory-maps the PLE safetensors, and uses an SSD prompt cache. MTP is fixed at depth 1. Prompt priming is capped at 8,192 tokens so a long request cannot retain a second full-context attention timeline. The context budget is total input plus generated output, so keep that total at or below 131,072 tokens.
+
+## Recommended generation settings
+
+Thinking mode is enabled by default. The upstream recommended settings are included in `generation_config.json`:
+
+```text
+temperature: 1.0
+top_p: 0.95
+top_k: 20
+min_p: 0.0
+presence_penalty: 0.0
+repetition_penalty: 1.0
+```
+
+For multi-turn agents, preserve `reasoning_content` between tool turns and keep `preserve_thinking` enabled.
+
+## Validation
+
+The artifact validator and strict lazy load found:
+
+- 22 expected and readable safetensors files
+- 2,916 indexed and stored tensors, including 62 MTP tensors
+- No missing, unindexed, duplicated, or misplaced tensors
+- No remaining per-expert MTP tensors after stacking
+- All 128 PLE mixed-bit overrides applied as intended
+- No missing or mismatched importance-matrix entries for measurable modules
+
+On an M5 Max with 128 GiB of unified memory, the SSD-backed PLE implementation loaded 64.65 GB of Metal-backed model arrays. A fresh request containing 131,071 uncached prompt tokens generated one token, for exactly 131,072 total tokens, in 873.41 seconds with oMLX's balanced memory guard enabled. The guard reclaimed 4.86 GB and 5.87 GB of pooled Metal buffers during prefill instead of weakening Apple's 107.52 GiB working-set ceiling.
+
+Live depth-1 requests activated the MTP path. The direct structured-call check accepted 39 of 45 draft tokens, or 86.7 percent. Across the six-turn Swival workflow, per-turn acceptance ranged from 72.5 to 87.7 percent. Acceptance is prompt-dependent and is not a quality guarantee.
+
+## Tool calling
+
+The original chat template and XML tool-call protocol are included. With MTP active, direct structured tool calls, continuation after a tool result, and a six-turn Swival file-write, checksum, read-back, and exact-byte comparison workflow were tested successfully. No identical tool invocation executed twice in that workflow.
+
+Tool calling is stochastic and is not guaranteed to be flawless. Synthetic exact-call testing passed 17 of 18 cases with resident PLE and 16 of 18 cases with SSD-backed PLE. One failed response emitted duplicate calls. Agents should reject duplicate identical calls before execution, allow at most one bounded repair attempt, and fail closed if the repaired response remains invalid.
+
+The dedicated MTP run passed 17 of 18 exact structured-call cases at the official thinking preset. The failed case emitted ordinary text claiming that it had called the requested tool but contained no structured call. The same case passed on an immediate isolated retry. MTP therefore matched the best prior non-MTP soak but did not eliminate the model's stochastic protocol failures.
+
+The duplicate-call protection used during validation is implemented by Swival and is not part of these model weights.
+
+## Limitations
+
+- This is a text-only conversion. Image and video inputs are unsupported.
+- MTP requires the bundled Qwen4 compatibility hook and oMLX's native-MTP scheduler. Stock MLX-LM cannot use this head.
+- Prompt priming is intentionally disabled for prompts above 8,192 tokens. Long prompts still run correctly but begin drafting without a primed head history.
+- The architecture supports 262,144 native tokens, but this conversion was validated and operationally limited to 131,072 total tokens.
+- Only the stated Apple Silicon and oMLX configuration was tested.
+- Quantization can reduce quality relative to the FP8 source.
+- Greedy MTP-on and MTP-off checks produced the same final JSON answer but not a bit-identical hidden reasoning sequence. Multi-token target verification can take numerically different matrix paths.
+- The model can produce incorrect content, malformed tool calls, duplicate tool calls, or unsafe actions. Tool execution requires independent validation and appropriate permission boundaries.
+- The model inherits the limitations and intended-use considerations of the upstream Qwen model.
+
+## Source and license
+
+The conversion was produced from revision `bcd9f01ddc9cff2316eb84281bebcd5b058bddce` of [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8).
+
+The weights remain subject to the Qwen Community License 1.0. Read `LICENSE` before using or redistributing the model, particularly its terms for commercial Model as a Service and AI Work Assistant businesses.
+
+The bundled `omlx_support` code is distributed under the Apache License 2.0 in `omlx_support/LICENSE` and includes modifications to the implementations proposed in oMLX pull requests #3161 and #3163.
+
+## Citation
+
+```bibtex
+@techreport{qwen2026design,
+    title       = {On the Design of {Qwen3.8-Next} Architecture: Evaluation, Efficiency, and Training Stability},
+    author      = {{Qwen Team}},
+    institution = {Alibaba Group},
+    month       = {August},
+    year        = {2026}
+}
+
+@misc{qwen3.8flashnext,
+    title  = {{Qwen3.8-Flash-Next}: A New Architecture, Towards Ultimate Cost-Efficiency},
+    author = {{Qwen Team}},
+    month  = {August},
+    year   = {2026},
+    url    = {https://qwen.ai/blog?id=qwen3.8-flash-next}
+}
+```

chat_template.jinja

@@ -0,0 +1,170 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- set reasoning_instructions = '' %}
+{%- if enable_thinking is undefined or enable_thinking is true %}
+    {%- set resolved_reasoning_effort = reasoning_effort|default('xhigh') %}
+    {%- if resolved_reasoning_effort not in ('xhigh', 'medium', 'low') %}
+        {{- raise_exception('Unexpected reasoning effort ' ~ reasoning_effort ~ '. Supported types are xhigh (default), medium, and low.') }}
+    {%- endif %}
+    {%- if resolved_reasoning_effort == 'xhigh' %}
+        {%- set reasoning_instructions = 'Reasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.' %}
+    {%- elif resolved_reasoning_effort == 'low' %}
+        {%- set reasoning_instructions = 'Reasoning effort is set to low. Keep your thinking brief and focused, moving directly to the conclusion without unnecessary elaboration.' %}
+    {%- endif %}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {%- if reasoning_instructions %}
+        {{- reasoning_instructions + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '<|im_start|>system\n' + (reasoning_instructions + '\n\n' if reasoning_instructions else '')  + content + '<|im_end|>\n' }}
+        {%- elif reasoning_instructions %}
+            {{- '<|im_start|>system\n' + reasoning_instructions + '<|im_end|>\n' }}
+        {%- endif %}
+    {%- elif reasoning_instructions %}
+        {{- '<|im_start|>system\n' + reasoning_instructions + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if preserve_thinking is undefined or preserve_thinking is true or loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined and tool_call.arguments != '' %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

generation_config.json

@@ -0,0 +1,12 @@
+{
+    "bos_token_id": 248044,
+    "do_sample": true,
+    "eos_token_id": [
+        248046,
+        248044
+    ],
+    "pad_token_id": 248044,
+    "temperature": 1.0,
+    "top_k": 20,
+    "top_p": 0.95
+}

omlx_support/LICENSE

@@ -0,0 +1,190 @@
+                                 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 the 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 the 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 any 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
+
+   Copyright 2025 oMLX contributors
+
+   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.

omlx_support/model_settings.json

@@ -0,0 +1,40 @@
+{
+  "version": 1,
+  "models": {
+    "Qwen3.8-Flash-Next-oQ4e-MTP-128k": {
+      "max_context_window": 131072,
+      "max_tokens": 32768,
+      "temperature": 1.0,
+      "top_p": 0.95,
+      "top_k": 20,
+      "min_p": 0.0,
+      "repetition_penalty": 1.0,
+      "presence_penalty": 0.0,
+      "force_sampling": true,
+      "chat_template_kwargs": {
+        "enable_thinking": true,
+        "preserve_thinking": true
+      },
+      "forced_ct_kwargs": [
+        "enable_thinking",
+        "preserve_thinking"
+      ],
+      "model_type_override": "llm",
+      "enable_thinking": true,
+      "preserve_thinking": true,
+      "reasoning_parser": "qwen",
+      "turboquant_kv_enabled": false,
+      "qwen35_ane_prefill_enabled": false,
+      "specprefill_enabled": false,
+      "dflash_enabled": false,
+      "mtp_enabled": true,
+      "mtp_num_draft_tokens": 1,
+      "vlm_mtp_enabled": false,
+      "is_pinned": false,
+      "is_default": true,
+      "trust_remote_code": false,
+      "display_name": "Qwen3.8 Flash Next oQ4e MTP 128k",
+      "description": "Text-only tool-use profile with depth-1 MTP and bounded prompt priming."
+    }
+  }
+}

omlx_support/qwen4_cache_integration.py

@@ -0,0 +1,381 @@
+# Qwen4 cache integration for oMLX.
+# Licensed under the Apache License 2.0. See LICENSE.
+
+from enum import Enum
+from typing import Any
+
+import mlx.core as mx
+
+from omlx.cache.type_handlers import (
+    CacheStateAxisInfo,
+    CacheTypeHandler,
+)
+from omlx.cache.type_registry import CacheTypeRegistry
+
+
+class Qwen4CacheType(Enum):
+    QSA_KV = "QSAKVCache"
+    QSA_QUANTIZED_KV = "QSAQuantizedKVCache"
+
+
+def _offset_from_meta(meta_state, fallback):
+    if isinstance(meta_state, (list, tuple)) and meta_state:
+        return int(meta_state[0])
+    if meta_state not in (None, ""):
+        return int(meta_state)
+    return fallback
+
+
+class QSAKVCacheHandler(CacheTypeHandler):
+    @property
+    def cache_type(self):
+        return Qwen4CacheType.QSA_KV
+
+    @property
+    def supports_block_slicing(self):
+        return True
+
+    def get_state_axis_info(self):
+        return (
+            CacheStateAxisInfo("keys", 2, True),
+            CacheStateAxisInfo("values", 2, True),
+            CacheStateAxisInfo("index_keys", 1, True),
+        )
+
+    def serialize_state(self, cache_obj):
+        keys, values, index_keys = cache_obj.state
+        return keys, values, index_keys
+
+    def serialize_meta_state(self, cache_obj):
+        return (int(cache_obj.offset),)
+
+    def extract_state(self, cache_obj):
+        elements = self.serialize_state(cache_obj)
+        return {
+            "keys": elements[0],
+            "values": elements[1],
+            "index_keys": elements[2],
+            "states": elements,
+            "cache_type": self.cache_type.value,
+        }
+
+    def get_seq_len(self, state):
+        keys = state.get("keys")
+        if keys is not None:
+            return int(keys.shape[2])
+        index_keys = state.get("index_keys")
+        return 0 if index_keys is None else int(index_keys.shape[1])
+
+    def slice_state(self, state, start_idx, end_idx):
+        keys = state.get("keys")
+        values = state.get("values")
+        index_keys = state.get("index_keys")
+        if keys is None or values is None:
+            return None
+        end_idx = min(end_idx, int(keys.shape[2]))
+        if start_idx >= end_idx:
+            return None
+        index_end = min(end_idx, int(index_keys.shape[1]))
+        elements = (
+            keys[:, :, start_idx:end_idx, :],
+            values[:, :, start_idx:end_idx, :],
+            index_keys[:, start_idx:index_end, :],
+        )
+        return {
+            "keys": elements[0],
+            "values": elements[1],
+            "index_keys": elements[2],
+            "states": elements,
+            "cache_type": self.cache_type.value,
+        }
+
+    def concatenate_states(self, states):
+        elements = [state.get("states") for state in states]
+        elements = [value for value in elements if value]
+        if not elements:
+            return {}
+        combined = (
+            mx.concatenate([value[0] for value in elements], axis=2),
+            mx.concatenate([value[1] for value in elements], axis=2),
+            mx.concatenate([value[2] for value in elements], axis=1),
+        )
+        return {
+            "keys": combined[0],
+            "values": combined[1],
+            "index_keys": combined[2],
+            "states": combined,
+            "cache_type": self.cache_type.value,
+        }
+
+    def deserialize_state(self, elements, meta_state=None):
+        from mlx_lm.models.qwen4_exp import QSAKVCache
+
+        keys = elements[0] if len(elements) > 0 else None
+        values = elements[1] if len(elements) > 1 else None
+        index_keys = elements[2] if len(elements) > 2 else None
+        fallback = 0 if keys is None else int(keys.shape[2])
+        cache = QSAKVCache()
+        cache.keys = keys
+        cache.values = values
+        cache.index_keys = index_keys
+        cache.offset = _offset_from_meta(meta_state, fallback)
+        return cache
+
+    def reconstruct_cache(self, state, meta_state=None):
+        elements = state.get("states")
+        if elements is None:
+            elements = (
+                state.get("keys"),
+                state.get("values"),
+                state.get("index_keys"),
+            )
+        return self.deserialize_state(tuple(elements), meta_state)
+
+
+class QSAQuantizedKVCacheHandler(CacheTypeHandler):
+    @property
+    def cache_type(self):
+        return Qwen4CacheType.QSA_QUANTIZED_KV
+
+    @property
+    def supports_block_slicing(self):
+        return True
+
+    def get_state_axis_info(self):
+        return (
+            CacheStateAxisInfo("key_weight", 2, True),
+            CacheStateAxisInfo("key_scales", 2, True),
+            CacheStateAxisInfo("key_biases", 2, True),
+            CacheStateAxisInfo("value_weight", 2, True),
+            CacheStateAxisInfo("value_scales", 2, True),
+            CacheStateAxisInfo("value_biases", 2, True),
+            CacheStateAxisInfo("index_keys", 1, True),
+        )
+
+    def serialize_state(self, cache_obj):
+        if cache_obj.keys is None:
+            return (None,) * 7
+        offset = int(cache_obj.offset)
+        keys = tuple(value[:, :, :offset, :] for value in cache_obj.keys)
+        values = tuple(value[:, :, :offset, :] for value in cache_obj.values)
+        index_keys = cache_obj.index_keys
+        if index_keys is not None:
+            index_keys = index_keys[:, :offset, :]
+        return (*keys, *values, index_keys)
+
+    def serialize_meta_state(self, cache_obj):
+        return (
+            int(cache_obj.offset),
+            int(cache_obj.group_size),
+            int(cache_obj.bits),
+        )
+
+    def extract_state(self, cache_obj):
+        elements = self.serialize_state(cache_obj)
+        return {
+            "states": elements,
+            "keys": elements[0],
+            "values": elements[3],
+            "index_keys": elements[6],
+            "cache_type": self.cache_type.value,
+        }
+
+    def get_seq_len(self, state):
+        keys = state.get("keys")
+        if keys is not None:
+            return int(keys.shape[2])
+        index_keys = state.get("index_keys")
+        return 0 if index_keys is None else int(index_keys.shape[1])
+
+    def slice_state(self, state, start_idx, end_idx):
+        elements = state.get("states")
+        if not elements or elements[0] is None:
+            return None
+        end_idx = min(end_idx, int(elements[0].shape[2]))
+        if start_idx >= end_idx:
+            return None
+        sliced = tuple(
+            value[:, :, start_idx:end_idx, :] for value in elements[:6]
+        ) + (elements[6][:, start_idx:end_idx, :],)
+        return {
+            "states": sliced,
+            "keys": sliced[0],
+            "values": sliced[3],
+            "index_keys": sliced[6],
+            "cache_type": self.cache_type.value,
+        }
+
+    def concatenate_states(self, states):
+        elements = [state.get("states") for state in states]
+        elements = [value for value in elements if value and value[0] is not None]
+        if not elements:
+            return {}
+        combined = tuple(
+            mx.concatenate([value[index] for value in elements], axis=2)
+            for index in range(6)
+        ) + (mx.concatenate([value[6] for value in elements], axis=1),)
+        return {
+            "states": combined,
+            "keys": combined[0],
+            "values": combined[3],
+            "index_keys": combined[6],
+            "cache_type": self.cache_type.value,
+        }
+
+    def deserialize_state(self, elements, meta_state=None):
+        from mlx_lm.models.qwen4_exp import QSAQuantizedKVCache
+
+        offset = _offset_from_meta(
+            meta_state,
+            0 if not elements or elements[0] is None else int(elements[0].shape[2]),
+        )
+        group_size = int(meta_state[1]) if meta_state and len(meta_state) > 1 else 64
+        bits = int(meta_state[2]) if meta_state and len(meta_state) > 2 else 4
+        cache = QSAQuantizedKVCache(group_size=group_size, bits=bits)
+        if elements and elements[0] is not None:
+            cache.keys = tuple(elements[:3])
+            cache.values = tuple(elements[3:6])
+        cache.index_keys = elements[6] if len(elements) > 6 else None
+        cache.offset = offset
+        return cache
+
+    def reconstruct_cache(self, state, meta_state=None):
+        return self.deserialize_state(tuple(state.get("states") or ()), meta_state)
+
+
+def _batch_indices(batch_indices):
+    if hasattr(batch_indices, "tolist"):
+        return [int(value) for value in batch_indices.tolist()]
+    return [int(value) for value in batch_indices]
+
+
+def _install_single_cache_batch_methods(cache_class):
+    def filter_rows(self, batch_indices):
+        indices = _batch_indices(batch_indices)
+        if not indices:
+            self.keys = None
+            self.values = None
+            self.index_keys = None
+            self.offset = 0
+            return
+        self.keys = _map_cache_arrays(self.keys, lambda value: value[indices])
+        self.values = _map_cache_arrays(self.values, lambda value: value[indices])
+        if self.index_keys is not None:
+            self.index_keys = self.index_keys[indices]
+
+    def extract_row(self, index):
+        result = type(self).__new__(type(self))
+        result.keys = _map_cache_arrays(self.keys, lambda value: value[index : index + 1])
+        result.values = _map_cache_arrays(self.values, lambda value: value[index : index + 1])
+        result.index_keys = (
+            None
+            if self.index_keys is None
+            else self.index_keys[index : index + 1]
+        )
+        result.offset = self.offset
+        if hasattr(self, "group_size"):
+            result.group_size = self.group_size
+            result.bits = self.bits
+        return result
+
+    def extend_rows(self, other):
+        if int(self.offset) != int(other.offset):
+            raise ValueError("QSA caches can only batch rows at the same offset")
+        if hasattr(self, "group_size") and (
+            self.group_size != other.group_size or self.bits != other.bits
+        ):
+            raise ValueError("quantized QSA caches must use the same layout")
+        offset = int(self.offset)
+        self.keys = _merge_cache_arrays(self.keys, other.keys, offset, 2)
+        self.values = _merge_cache_arrays(self.values, other.values, offset, 2)
+        self.index_keys = _merge_cache_arrays(
+            self.index_keys,
+            other.index_keys,
+            offset,
+            1,
+        )
+
+    @classmethod
+    def merge_rows(cls, caches):
+        caches = list(caches)
+        if not caches:
+            return cls()
+        result = caches[0].extract(0)
+        for cache in caches[1:]:
+            result.extend(cache)
+        return result
+
+    cache_class.filter = filter_rows
+    cache_class.extract = extract_row
+    cache_class.extend = extend_rows
+    cache_class.merge = merge_rows
+
+
+def _install_quantized_state_layout(cache_class):
+    def get_state(self):
+        if self.keys is None:
+            return (None,) * 7
+        offset = int(self.offset)
+        keys = tuple(value[:, :, :offset, :] for value in self.keys)
+        values = tuple(value[:, :, :offset, :] for value in self.values)
+        index_keys = self.index_keys
+        if index_keys is not None:
+            index_keys = index_keys[:, :offset, :]
+        return (*keys, *values, index_keys)
+
+    def set_state(self, value):
+        if len(value) == 2 and isinstance(value[0], (list, tuple)):
+            quantized_state, self.index_keys = value
+            self.keys, self.values = quantized_state
+        else:
+            self.keys = tuple(value[:3]) if value and value[0] is not None else None
+            self.values = tuple(value[3:6]) if value and value[3] is not None else None
+            self.index_keys = value[6] if len(value) > 6 else None
+        self.offset = 0 if self.keys is None else int(self.keys[0].shape[2])
+
+    cache_class.state = property(get_state, set_state)
+
+
+def _map_cache_arrays(value, function):
+    if value is None:
+        return None
+    if isinstance(value, (list, tuple)):
+        return tuple(function(item) for item in value)
+    return function(value)
+
+
+def _merge_cache_arrays(left, right, length, sequence_axis):
+    if left is None:
+        return right
+    if right is None:
+        return left
+    if isinstance(left, (list, tuple)):
+        return tuple(
+            _merge_cache_arrays(a, b, length, sequence_axis)
+            for a, b in zip(left, right)
+        )
+    slices = [slice(None)] * left.ndim
+    slices[sequence_axis] = slice(0, length)
+    slices = tuple(slices)
+    return mx.concatenate([left[slices], right[slices]], axis=0)
+
+
+def register_qwen4_cache_integration():
+    import mlx_lm.models.cache as mlx_cache
+
+    from mlx_lm.models.qwen4_exp import QSAKVCache, QSAQuantizedKVCache
+
+    handlers = (QSAKVCacheHandler(), QSAQuantizedKVCacheHandler())
+    for handler in handlers:
+        CacheTypeRegistry.register(handler)
+    CacheTypeRegistry._class_name_map.update(
+        {
+            QSAKVCache.__name__: Qwen4CacheType.QSA_KV,
+            QSAQuantizedKVCache.__name__: Qwen4CacheType.QSA_QUANTIZED_KV,
+        }
+    )
+    _install_single_cache_batch_methods(QSAKVCache)
+    _install_single_cache_batch_methods(QSAQuantizedKVCache)
+    _install_quantized_state_layout(QSAQuantizedKVCache)
+    mlx_cache.QSAKVCache = QSAKVCache
+    mlx_cache.QSAQuantizedKVCache = QSAQuantizedKVCache

omlx_support/qwen4_exp.py

@@ -0,0 +1,1361 @@
+from __future__ import annotations
+
+import json
+import logging
+import math
+import mmap
+import os
+import struct
+from copy import copy
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from mlx_lm.models.activations import swiglu
+from mlx_lm.models.base import BaseModelArgs, create_ssm_mask
+from mlx_lm.models.cache import ArraysCache, KVCache, QuantizedKVCache
+from mlx_lm.models.gated_delta import gated_delta_update
+from mlx_lm.models.switch_layers import SwitchGLU
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "qwen4_exp"
+    vocab_size: int = 248320
+    hidden_size: int = 2048
+    num_hidden_layers: int = 40
+    mtp_num_hidden_layers: int = 0
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 2
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 32768
+    rms_norm_eps: float = 1e-6
+    tie_word_embeddings: bool = False
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    head_dim: int = 256
+    linear_conv_kernel_dim: int = 4
+    linear_key_head_dim: int = 128
+    linear_value_head_dim: int = 128
+    linear_num_key_heads: int = 16
+    linear_num_value_heads: int = 32
+    moe_intermediate_size: int = 512
+    shared_expert_intermediate_size: int = 512
+    num_experts_per_tok: int = 10
+    num_experts: int = 512
+    layer_types: List[str] = field(default_factory=list)
+    hc_count: int = 4
+    hc_lowrank: int = 320
+    ple_layer_ids: List[int] = field(default_factory=list)
+    ple_embed_dim: Optional[int] = None
+    ple_conv_kernel_size: int = 4
+    ngram_size: int = 3
+    heads_per_ngram: int = 8
+    ngram_vocab_size_base: int = 20_000_000
+    make_ngram_vocab_size_divisible_by: int = 128
+    seed: int = 1234
+    split_ngram_parts: int = 512
+    indexer_n_heads: Optional[int] = None
+    indexer_kv_heads: Optional[int] = None
+    indexer_head_dim: Optional[int] = None
+    indexer_budget: Optional[int] = None
+    indexer_compress_ratio: Optional[int] = None
+    norm_topk_prob: bool = True
+    output_gate_type: Optional[str] = None
+    eos_token_id: Optional[int | List[int]] = None
+    rope_theta: float = 10000.0
+    partial_rotary_factor: float = 1.0
+    rope_parameters: Optional[Dict[str, Any]] = None
+
+    @classmethod
+    def from_dict(cls, params):
+        source = dict(params.get("text_config", params))
+        source["model_type"] = params.get("model_type", source.get("model_type", "qwen4_exp"))
+        if params.get("eos_token_id") is not None:
+            source["eos_token_id"] = params["eos_token_id"]
+        rope = source.get("rope_parameters") or {}
+        source.setdefault("rope_theta", rope.get("rope_theta", 10000.0))
+        source.setdefault("partial_rotary_factor", rope.get("partial_rotary_factor", 1.0))
+        return super().from_dict(source)
+
+    def __post_init__(self):
+        if self.ple_embed_dim is None:
+            self.ple_embed_dim = self.hidden_size
+        if not self.layer_types:
+            self.layer_types = [
+                "linear_attention" if (i + 1) % 4 else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        if len(self.layer_types) != self.num_hidden_layers:
+            raise ValueError("layer_types must contain one entry per hidden layer")
+        if self.hc_count <= 1:
+            raise ValueError("hc_count must be greater than one")
+        if self.linear_num_value_heads % self.linear_num_key_heads:
+            raise ValueError("linear value heads must be divisible by key heads")
+        if not 0 < self.num_experts_per_tok <= self.num_experts:
+            raise ValueError("num_experts_per_tok must select existing experts")
+        if any(layer < 1 or layer > self.num_hidden_layers for layer in self.ple_layer_ids):
+            raise ValueError("PLE layer ids are one-indexed hidden-layer ids")
+        if self.ple_layer_ids and self.eos_token_id is None:
+            raise ValueError("PLE requires eos_token_id")
+        qsa = (
+            self.indexer_n_heads,
+            self.indexer_kv_heads,
+            self.indexer_head_dim,
+            self.indexer_budget,
+            self.indexer_compress_ratio,
+        )
+        if any(value is None for value in qsa):
+            raise ValueError("QSA requires every indexer field")
+        if self.indexer_kv_heads != 1:
+            raise ValueError("QSA requires one indexer key head")
+        if self.indexer_budget % self.indexer_compress_ratio:
+            raise ValueError("indexer_budget must divide into complete compressed blocks")
+        ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
+        if self.ple_layer_ids and self.ple_embed_dim % ngram_heads:
+            raise ValueError("ple_embed_dim must be divisible by its n-gram heads")
+
+
+class Qwen4RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float, group_size: Optional[int] = None):
+        super().__init__()
+        self.weight = mx.zeros(dim)
+        self.eps = eps
+        self.group_size = group_size
+
+    def __call__(self, x):
+        dtype = x.dtype
+        value = x.astype(mx.float32)
+        if self.group_size is not None:
+            value = value.reshape(*value.shape[:-1], -1, self.group_size)
+        value = value * mx.rsqrt(mx.mean(mx.square(value), axis=-1, keepdims=True) + self.eps)
+        if self.group_size is not None:
+            value = value.reshape(*x.shape)
+        return (value * (1.0 + self.weight.astype(mx.float32))).astype(dtype)
+
+
+class Qwen4RMSNormGated(nn.Module):
+    def __init__(self, dim: int, eps: float, activation: str):
+        super().__init__()
+        self.weight = mx.ones(dim)
+        self.eps = eps
+        self.activation = activation
+
+    def __call__(self, x, gate):
+        dtype = x.dtype
+        value = mx.fast.rms_norm(x, self.weight, self.eps)
+        gate = gate.astype(mx.float32)
+        gate = mx.sigmoid(gate) if self.activation == "sigmoid" else nn.silu(gate)
+        return (value.astype(mx.float32) * gate).astype(dtype)
+
+
+def _l2_normalize(x, eps=1e-6):
+    return x * mx.rsqrt(mx.sum(x * x, axis=-1, keepdims=True) + eps)
+
+
+def _apply_rope(x, positions, rotary_dim: int, theta: float):
+    if rotary_dim == 0:
+        return x
+    dtype = x.dtype
+    inv_freq = theta ** (-mx.arange(0, rotary_dim, 2, dtype=mx.float32) / rotary_dim)
+    angles = positions.astype(mx.float32)[..., None] * inv_freq
+    cos = mx.concatenate([mx.cos(angles), mx.cos(angles)], axis=-1)
+    sin = mx.concatenate([mx.sin(angles), mx.sin(angles)], axis=-1)
+    while cos.ndim < x.ndim:
+        cos = mx.expand_dims(cos, axis=-2)
+        sin = mx.expand_dims(sin, axis=-2)
+    rotated, remainder = x[..., :rotary_dim], x[..., rotary_dim:]
+    first, second = mx.split(rotated, 2, axis=-1)
+    rotated = rotated * cos + mx.concatenate([-second, first], axis=-1) * sin
+    return mx.concatenate([rotated, remainder], axis=-1).astype(dtype)
+
+
+class QSAKVCache(KVCache):
+    def __init__(self):
+        super().__init__()
+        self.index_keys = None
+
+    def update_indexer(self, keys):
+        previous = self.offset
+        length = keys.shape[1]
+        end = previous + length
+        if self.index_keys is None or end > self.index_keys.shape[1]:
+            batch, _, head_dim = keys.shape
+            steps = (self.step + length - 1) // self.step
+            extension = mx.zeros(
+                (batch, steps * self.step, head_dim),
+                dtype=keys.dtype,
+            )
+            if self.index_keys is None:
+                self.index_keys = extension
+            else:
+                if previous % self.step:
+                    self.index_keys = self.index_keys[:, :previous, :]
+                self.index_keys = mx.concatenate([self.index_keys, extension], axis=1)
+        self.index_keys[:, previous:end, :] = keys
+        return self.index_keys[:, :end, :]
+
+    @property
+    def state(self):
+        index_state = (
+            self.index_keys
+            if self.index_keys is None
+            else self.index_keys[:, : self.offset, :]
+        )
+        if self.keys is None:
+            return self.keys, self.values, index_state
+        return (
+            self.keys[..., : self.offset, :],
+            self.values[..., : self.offset, :],
+            index_state,
+        )
+
+    @state.setter
+    def state(self, value):
+        self.keys, self.values, self.index_keys = value
+        self.offset = 0 if self.keys is None else self.keys.shape[2]
+
+    def trim(self, count):
+        return super().trim(count)
+
+    def to_quantized(self, group_size=64, bits=4):
+        cache = QSAQuantizedKVCache(group_size=group_size, bits=bits)
+        cache.offset = self.offset
+        cache.index_keys = self.index_keys
+        if self.keys is not None:
+            cache.keys = mx.quantize(self.keys[..., : self.offset, :], group_size=group_size, bits=bits)
+            cache.values = mx.quantize(self.values[..., : self.offset, :], group_size=group_size, bits=bits)
+        return cache
+
+    @property
+    def nbytes(self):
+        size = super().nbytes
+        return size + (0 if self.index_keys is None else self.index_keys.nbytes)
+
+
+class QSAQuantizedKVCache(QuantizedKVCache):
+    def __init__(self, group_size=64, bits=4):
+        super().__init__(group_size=group_size, bits=bits)
+        self.index_keys = None
+
+    def update_indexer(self, keys):
+        previous = self.offset
+        length = keys.shape[1]
+        end = previous + length
+        if self.index_keys is None or end > self.index_keys.shape[1]:
+            batch, _, head_dim = keys.shape
+            steps = (self.step + length - 1) // self.step
+            extension = mx.zeros(
+                (batch, steps * self.step, head_dim),
+                dtype=keys.dtype,
+            )
+            if self.index_keys is None:
+                self.index_keys = extension
+            else:
+                if previous % self.step:
+                    self.index_keys = self.index_keys[:, :previous, :]
+                self.index_keys = mx.concatenate([self.index_keys, extension], axis=1)
+        self.index_keys[:, previous:end, :] = keys
+        return self.index_keys[:, :end, :]
+
+    @property
+    def state(self):
+        quantized_state = (self.keys, self.values) if self.keys is None else super().state
+        index_state = (
+            self.index_keys
+            if self.index_keys is None
+            else self.index_keys[:, : self.offset, :]
+        )
+        return quantized_state, index_state
+
+    @state.setter
+    def state(self, value):
+        quantized_state, self.index_keys = value
+        QuantizedKVCache.state.fset(self, quantized_state)
+        self.offset = 0 if self.keys is None else self.keys[0].shape[2]
+
+    def trim(self, count):
+        return super().trim(count)
+
+    def to_quantized(self, group_size=64, bits=4):
+        if group_size == self.group_size and bits == self.bits:
+            return self
+        raise ValueError("QSA KV cache is already quantized")
+
+    @property
+    def nbytes(self):
+        size = super().nbytes
+        return size + (0 if self.index_keys is None else self.index_keys.nbytes)
+
+
+class Qwen4GatedDeltaNet(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_v_heads = args.linear_num_value_heads
+        self.num_k_heads = args.linear_num_key_heads
+        self.head_k_dim = args.linear_key_head_dim
+        self.head_v_dim = args.linear_value_head_dim
+        self.key_dim = self.num_k_heads * self.head_k_dim
+        self.value_dim = self.num_v_heads * self.head_v_dim
+        self.conv_kernel_size = args.linear_conv_kernel_dim
+        self.conv_dim = 2 * self.key_dim + self.value_dim
+        self.conv1d = nn.Conv1d(
+            self.conv_dim,
+            self.conv_dim,
+            self.conv_kernel_size,
+            groups=self.conv_dim,
+            bias=False,
+        )
+        self.in_proj_qkv = nn.Linear(args.hidden_size, self.conv_dim, bias=False)
+        self.in_proj_z = nn.Linear(args.hidden_size, self.value_dim, bias=False)
+        self.in_proj_b = nn.Linear(args.hidden_size, self.num_v_heads, bias=False)
+        self.in_proj_a = nn.Linear(args.hidden_size, self.num_v_heads, bias=False)
+        self.dt_bias = mx.ones(self.num_v_heads)
+        self.A_log = mx.log(mx.random.uniform(low=0.01, high=16.0, shape=(self.num_v_heads,)))
+        self.norm = Qwen4RMSNormGated(
+            self.head_v_dim,
+            args.rms_norm_eps,
+            args.output_gate_type or args.hidden_act,
+        )
+        self.out_proj = nn.Linear(self.value_dim, args.hidden_size, bias=False)
+
+    def _process_chunk(self, mixed, a, b, conv_state, state, mask=None):
+        batch, length = mixed.shape[:2]
+        conv_input = mx.concatenate([conv_state, mixed], axis=1)
+        next_conv_state = mx.contiguous(
+            conv_input[:, -(self.conv_kernel_size - 1) :, :]
+        )
+        mixed = nn.silu(self.conv1d(conv_input))
+        q, k, value = mx.split(mixed, [self.key_dim, 2 * self.key_dim], axis=-1)
+        q = q.reshape(batch, length, self.num_k_heads, self.head_k_dim)
+        k = k.reshape(batch, length, self.num_k_heads, self.head_k_dim)
+        value = value.reshape(batch, length, self.num_v_heads, self.head_v_dim)
+        q = _l2_normalize(q) * (self.head_k_dim**-0.5)
+        k = _l2_normalize(k)
+        output, state = gated_delta_update(
+            q,
+            k,
+            value,
+            a.astype(mx.float32),
+            b,
+            self.A_log,
+            self.dt_bias,
+            state,
+            mask,
+            use_kernel=not self.training,
+        )
+        return output, next_conv_state, state
+
+    def __call__(self, x, mask=None, cache=None, n_confirmed=0):
+        batch, length, _ = x.shape
+        mixed = self.in_proj_qkv(x)
+        z = self.in_proj_z(x).reshape(batch, length, self.num_v_heads, self.head_v_dim)
+        b = self.in_proj_b(x)
+        a = self.in_proj_a(x)
+        if mask is not None:
+            mixed = mx.where(mask[..., None], mixed, 0)
+        if cache is not None and cache[0] is not None:
+            conv_state = cache[0]
+        else:
+            conv_state = mx.zeros((batch, self.conv_kernel_size - 1, self.conv_dim), dtype=x.dtype)
+        state = None if cache is None else cache[1]
+        output, next_conv_state, next_state = self._process_chunk(
+            mixed,
+            a,
+            b,
+            conv_state,
+            state,
+            mask,
+        )
+        if cache is not None:
+            cache[0] = next_conv_state
+            cache[1] = next_state
+            if 0 < n_confirmed < length:
+                cache.rollback_state = (conv_state, state)
+                cache._mtp_draft_stash = (mixed, a, b, mask)
+            cache.advance(length)
+        output = self.norm(output, z).reshape(batch, length, -1)
+        return self.out_proj(output)
+
+
+class Qwen4QSAIndexer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.n_heads = args.indexer_n_heads
+        self.head_dim = args.indexer_head_dim
+        self.budget = args.indexer_budget
+        self.compress_ratio = args.indexer_compress_ratio
+        self.block_topk = self.budget // self.compress_ratio
+        self.rotary_dim = int(args.head_dim * args.partial_rotary_factor)
+        self.rope_theta = args.rope_theta
+        self.index_qk_proj = nn.Linear(
+            args.hidden_size,
+            (args.indexer_n_heads + args.indexer_kv_heads) * args.indexer_head_dim,
+            bias=False,
+        )
+        self.q_layernorm = Qwen4RMSNorm(self.head_dim, args.rms_norm_eps)
+        self.k_layernorm = Qwen4RMSNorm(self.head_dim, args.rms_norm_eps)
+
+    def __call__(self, hidden_states, cache, offset):
+        batch, length, _ = hidden_states.shape
+        qk = self.index_qk_proj(hidden_states)
+        split = self.n_heads * self.head_dim
+        query, raw_keys = mx.split(qk, [split], axis=-1)
+        query = self.q_layernorm(query.reshape(batch, length, self.n_heads, self.head_dim))
+        raw_keys = raw_keys.reshape(batch, length, self.head_dim)
+        raw_keys = cache.update_indexer(raw_keys) if cache is not None else raw_keys
+        positions = offset + mx.arange(length)
+        query = _apply_rope(query, positions[None], self.rotary_dim, self.rope_theta)
+        return query, raw_keys, positions
+
+    def select(self, query, raw_keys, positions):
+        batch, length, _, _ = query.shape
+        key_length = raw_keys.shape[1]
+        ratio = self.compress_ratio
+        block_count = key_length // ratio
+        selected_parts = []
+        valid_parts = []
+        if block_count:
+            pooled = raw_keys[:, : block_count * ratio].reshape(
+                batch, block_count, ratio, self.head_dim
+            ).mean(axis=2)
+            pooled = self.k_layernorm(pooled)
+            block_positions = mx.arange(block_count) * ratio
+            pooled = _apply_rope(
+                pooled,
+                block_positions[None],
+                self.rotary_dim,
+                self.rope_theta,
+            )
+            scores = mx.einsum(
+                "blhd,bkd->blhk",
+                query.astype(mx.float32),
+                pooled.astype(mx.float32),
+            )
+            scores = mx.sum(mx.maximum(scores, 0), axis=2) / math.sqrt(self.head_dim)
+            complete = (positions + 1) // ratio
+            block_valid = mx.arange(block_count)[None, None, :] < complete[None, :, None]
+            scores = mx.where(block_valid, scores, mx.finfo(scores.dtype).min)
+            take = min(self.block_topk, block_count)
+            if take == block_count:
+                selected_blocks = mx.broadcast_to(
+                    mx.arange(block_count)[None, None, :],
+                    (batch, length, block_count),
+                )
+            else:
+                selected_blocks = mx.argpartition(scores, kth=block_count - take, axis=-1)[..., -take:]
+            chosen_valid = mx.take_along_axis(block_valid, selected_blocks, axis=-1)
+            selected_parts.append(
+                (selected_blocks[..., None] * ratio + mx.arange(ratio)).reshape(batch, length, -1)
+            )
+            valid_parts.append(
+                mx.broadcast_to(chosen_valid[..., None], (*chosen_valid.shape, ratio)).reshape(batch, length, -1)
+            )
+        tail_width = max(ratio - 1, 1)
+        tail_start = ((positions + 1) // ratio) * ratio
+        tail = tail_start[:, None] + mx.arange(tail_width)[None]
+        tail_valid = tail <= positions[:, None]
+        selected_parts.append(mx.broadcast_to(tail[None], (batch, length, tail_width)))
+        valid_parts.append(mx.broadcast_to(tail_valid[None], (batch, length, tail_width)))
+        return mx.concatenate(selected_parts, axis=-1), mx.concatenate(valid_parts, axis=-1)
+
+
+class Qwen4Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_heads = args.num_attention_heads
+        self.num_kv_heads = args.num_key_value_heads
+        self.head_dim = args.head_dim
+        self.repeats = self.num_heads // self.num_kv_heads
+        self.scale = self.head_dim**-0.5
+        self.rotary_dim = int(self.head_dim * args.partial_rotary_factor)
+        self.rope_theta = args.rope_theta
+        self.q_proj = nn.Linear(args.hidden_size, self.num_heads * self.head_dim * 2, bias=args.attention_bias)
+        self.k_proj = nn.Linear(args.hidden_size, self.num_kv_heads * self.head_dim, bias=args.attention_bias)
+        self.v_proj = nn.Linear(args.hidden_size, self.num_kv_heads * self.head_dim, bias=args.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, args.hidden_size, bias=args.attention_bias)
+        self.q_norm = Qwen4RMSNorm(self.head_dim, args.rms_norm_eps)
+        self.k_norm = Qwen4RMSNorm(self.head_dim, args.rms_norm_eps)
+        self.indexer = Qwen4QSAIndexer(args)
+
+    def _select_cache_rows(self, values, batch_index, indices, cache):
+        if not isinstance(values, (list, tuple)):
+            return values[batch_index, :, indices, :].transpose(2, 0, 1, 3)
+        parts = [value[batch_index, :, indices, :].transpose(2, 0, 1, 3) for value in values]
+        return mx.dequantize(
+            parts[0],
+            parts[1],
+            parts[2],
+            group_size=cache.group_size,
+            bits=cache.bits,
+        )
+
+    def _sparse_attention(self, query, keys, values, selected, valid, cache):
+        outputs = []
+        for batch_index in range(query.shape[0]):
+            grouped_query = query[batch_index].reshape(
+                self.num_kv_heads, self.repeats, query.shape[2], self.head_dim
+            )
+            chunks = []
+            for start in range(0, query.shape[2], 64):
+                end = min(start + 64, query.shape[2])
+                key_length = (
+                    keys[0].shape[2]
+                    if isinstance(keys, (list, tuple))
+                    else keys.shape[2]
+                )
+                indices = mx.clip(
+                    selected[batch_index, start:end],
+                    0,
+                    key_length - 1,
+                )
+                selected_keys = self._select_cache_rows(keys, batch_index, indices, cache)
+                selected_values = self._select_cache_rows(values, batch_index, indices, cache)
+                local_query = grouped_query[:, :, start:end]
+                scores = mx.einsum("hrld,hlmd->hrlm", local_query, selected_keys) * self.scale
+                scores = mx.where(valid[batch_index, start:end][None, None], scores, mx.finfo(scores.dtype).min)
+                probabilities = mx.softmax(scores, axis=-1, precise=True)
+                output = mx.einsum("hrlm,hlmd->hrld", probabilities, selected_values)
+                mx.eval(output)
+                chunks.append(output)
+            output = mx.concatenate(chunks, axis=2)
+            outputs.append(output.reshape(self.num_heads, query.shape[2], self.head_dim)[None])
+        return mx.concatenate(outputs, axis=0)
+
+    def __call__(self, x, cache=None):
+        batch, length, _ = x.shape
+        offset = 0 if cache is None else cache.offset
+        projected = self.q_proj(x).reshape(batch, length, self.num_heads, 2 * self.head_dim)
+        query, gate = mx.split(projected, 2, axis=-1)
+        gate = gate.reshape(batch, length, -1)
+        keys = self.k_proj(x).reshape(batch, length, self.num_kv_heads, self.head_dim)
+        values = self.v_proj(x).reshape(batch, length, self.num_kv_heads, self.head_dim)
+        positions = offset + mx.arange(length)
+        query = _apply_rope(self.q_norm(query), positions[None], self.rotary_dim, self.rope_theta)
+        keys = _apply_rope(self.k_norm(keys), positions[None], self.rotary_dim, self.rope_theta)
+        query = query.transpose(0, 2, 1, 3)
+        keys = keys.transpose(0, 2, 1, 3)
+        values = values.transpose(0, 2, 1, 3)
+        index_query, raw_keys, _ = self.indexer(x, cache, offset)
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+        selected, valid = self.indexer.select(index_query, raw_keys, positions)
+        output = self._sparse_attention(query, keys, values, selected, valid, cache)
+        output = output.transpose(0, 2, 1, 3).reshape(batch, length, -1)
+        return self.o_proj(output * mx.sigmoid(gate))
+
+
+class Qwen4MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+
+    def __call__(self, x):
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+
+
+class Qwen4SparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.top_k = args.num_experts_per_tok
+        self.norm_topk_prob = args.norm_topk_prob
+        self.gate = nn.Linear(args.hidden_size, args.num_experts, bias=False)
+        self.switch_mlp = SwitchGLU(args.hidden_size, args.moe_intermediate_size, args.num_experts)
+        self.shared_expert = Qwen4MLP(args.hidden_size, args.shared_expert_intermediate_size)
+        self.shared_expert_gate = nn.Linear(args.hidden_size, 1, bias=False)
+
+    def _routing_weights(self, router_logits):
+        gates = mx.softmax(router_logits.astype(mx.float32), axis=-1, precise=True)
+        indices = mx.argpartition(gates, kth=-self.top_k, axis=-1)[..., -self.top_k :]
+        scores = mx.take_along_axis(gates, indices, axis=-1)
+        if self.norm_topk_prob:
+            scores = scores / mx.sum(scores, axis=-1, keepdims=True)
+        return scores.astype(router_logits.dtype), indices
+
+    def __call__(self, x):
+        scores, indices = self._routing_weights(self.gate(x))
+        routed = mx.sum(self.switch_mlp(x, indices) * scores[..., None], axis=-2)
+        shared = mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
+        return routed + shared
+
+
+class Qwen4GatedResidual(nn.Module):
+    def __init__(self, args: ModelArgs, combine=True):
+        super().__init__()
+        self.hc_count = args.hc_count
+        self.hidden_size = args.hidden_size
+        total = self.hc_count * self.hidden_size
+        self.hc_norm = Qwen4RMSNorm(total, args.rms_norm_eps, group_size=self.hidden_size)
+        self.input_mix_weight_down = nn.Linear(total, args.hc_lowrank, bias=False)
+        self.input_mix_weight_up = nn.Linear(args.hc_lowrank, total, bias=False)
+        self.block_inject_weight = nn.Linear(total, self.hc_count, bias=False) if combine else None
+
+    def __call__(self, x):
+        normalized = self.hc_norm(x)
+        weights = nn.silu(self.input_mix_weight_down(normalized) / self.hc_count)
+        weights = mx.sigmoid(self.input_mix_weight_up(weights)).reshape(
+            *x.shape[:-1], self.hc_count, self.hidden_size
+        )
+        mixed = mx.mean(
+            weights * normalized.reshape(*x.shape[:-1], self.hc_count, self.hidden_size),
+            axis=-2,
+        )
+        if self.block_inject_weight is None:
+            return mixed
+        injection = 2 * mx.sigmoid(self.block_inject_weight(normalized) / self.hc_count)
+        return mixed, x, injection
+
+
+_MASK64 = (1 << 64) - 1
+_SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
+_SPLITMIX_M1 = 0xBF58476D1CE4E5B9
+_SPLITMIX_M2 = 0x94D049BB133111EB
+
+
+def _splitmix64(value):
+    value = (value + _SPLITMIX_GAMMA) & _MASK64
+    value = ((value ^ (value >> 30)) * _SPLITMIX_M1) & _MASK64
+    value = ((value ^ (value >> 27)) * _SPLITMIX_M2) & _MASK64
+    return (value ^ (value >> 31)) & _MASK64
+
+
+def _multipliers(vocab_size, ngram_size, layer_index, seed):
+    maximum = ((1 << 63) - 1) // max(vocab_size, 1)
+    bound = max(1, maximum // 2)
+    base = seed + 10007 * layer_index
+    return [
+        2 * (_splitmix64((base + _SPLITMIX_GAMMA * (index + 1)) & _MASK64) % bound) + 1
+        for index in range(ngram_size)
+    ]
+
+
+def _is_prime(value):
+    if value < 2:
+        return False
+    if value % 2 == 0:
+        return value == 2
+    return all(value % divisor for divisor in range(3, math.isqrt(value) + 1, 2))
+
+
+def _nth_prime_after(start, count):
+    value = start
+    for _ in range(count):
+        value += 1
+        while not _is_prime(value):
+            value += 1
+    return value
+
+
+class Qwen4NGramEmbedding(nn.Module):
+    def __init__(self, args: ModelArgs, layer_index: int, model_layer_index=None):
+        super().__init__()
+        self.ngram_size = args.ngram_size
+        self.context_len = self.ngram_size - 1
+        self.heads_per_ngram = args.heads_per_ngram
+        self.ngram_heads = self.context_len * self.heads_per_ngram
+        self.eos_token_id = args.eos_token_id[0] if isinstance(args.eos_token_id, list) else args.eos_token_id
+        sizes = [
+            _nth_prime_after(args.ngram_vocab_size_base - 1, head + 1)
+            for head in range(self.ngram_heads)
+        ]
+        offsets = []
+        total = 0
+        for size in sizes:
+            offsets.append(total)
+            total += size
+        padded = math.ceil(total / args.make_ngram_vocab_size_divisible_by) * args.make_ngram_vocab_size_divisible_by
+        self.layer_multipliers = mx.array(
+            _multipliers(args.vocab_size, self.ngram_size, layer_index, args.seed),
+            dtype=mx.int64,
+        )
+        self.ngram_heads_vocab_sizes = mx.array(sizes, dtype=mx.int64)
+        self.ngram_heads_offsets = mx.array(offsets, dtype=mx.int64)
+        model_path = os.environ.get("OMLX_QWEN4_PLE_MODEL_PATH")
+        mode = os.environ.get("OMLX_QWEN4_PLE_MODE", "resident")
+        embedding_args = (
+            padded,
+            args.ple_embed_dim // self.ngram_heads,
+            args.split_ngram_parts,
+        )
+        if mode == "mmap":
+            if model_path is None or model_layer_index is None:
+                raise RuntimeError("SSD-backed PLE requires its model path and decoder layer index")
+            prefix = f"model.layers.{model_layer_index}.ple.ple_embedding.ngram_embedding"
+            self.ngram_embedding = DiskBackedShardedEmbedding(
+                model_path,
+                prefix,
+                *embedding_args,
+            )
+        elif mode == "resident":
+            self.ngram_embedding = ShardedEmbedding(*embedding_args)
+        else:
+            raise ValueError("OMLX_QWEN4_PLE_MODE must be resident or mmap")
+
+    def _shift(self, tokens, shift):
+        if shift == 0:
+            return tokens
+        batch, length = tokens.shape
+        positions = mx.arange(length)
+        eos_positions = mx.where(tokens == self.eos_token_id, positions[None], -1)
+        inclusive = mx.cummax(eos_positions, axis=1)
+        previous = mx.concatenate([mx.full((batch, 1), -1, dtype=mx.int64), inclusive[:, :-1]], axis=1)
+        source = positions - shift
+        gathered = tokens[:, mx.maximum(source, 0)]
+        valid = (positions[None] - previous - 1 >= shift) & (source[None] >= 0)
+        return mx.where(valid, gathered, self.eos_token_id)
+
+    def __call__(self, input_ids, cache=None):
+        input_ids = input_ids.astype(mx.int64)
+        if cache is not None and cache[3] is not None:
+            previous = cache[3]
+        else:
+            previous = mx.full((input_ids.shape[0], self.context_len), self.eos_token_id, dtype=mx.int64)
+        history = mx.concatenate([previous, input_ids], axis=1)
+        if cache is not None:
+            cache[3] = mx.contiguous(history[:, -self.context_len :])
+        shifted = [self._shift(history, shift) for shift in range(self.ngram_size)]
+        blocks = []
+        for ngram in range(2, self.ngram_size + 1):
+            start = (ngram - 2) * self.heads_per_ngram
+            end = start + self.heads_per_ngram
+            mixed = shifted[0] * self.layer_multipliers[0]
+            for position in range(1, ngram):
+                mixed = mx.bitwise_xor(mixed, shifted[position] * self.layer_multipliers[position])
+            ids = mixed[..., None] % self.ngram_heads_vocab_sizes[start:end]
+            blocks.append(ids + self.ngram_heads_offsets[start:end])
+        ids = mx.concatenate(blocks, axis=-1)[:, -input_ids.shape[1] :]
+        return self.ngram_embedding(ids).reshape(input_ids.shape[0], input_ids.shape[1], -1)
+
+
+class ShardedEmbedding(nn.Module):
+    def __init__(self, num_embeddings, dims, shard_count):
+        super().__init__()
+        if num_embeddings % shard_count:
+            raise ValueError("the padded n-gram vocabulary must divide evenly into shards")
+        self.shard_size = num_embeddings // shard_count
+        self.dims = dims
+        self.shards = [nn.Embedding(self.shard_size, dims) for _ in range(shard_count)]
+
+    def __call__(self, ids):
+        shape = ids.shape
+        flat_ids = ids.reshape(-1)
+        if flat_ids.size == 0:
+            return self.shards[0](flat_ids).reshape(*shape, self.dims)
+        shard_ids = flat_ids // self.shard_size
+        local_ids = flat_ids % self.shard_size
+        mx.eval(shard_ids)
+        host_shards = shard_ids.tolist()
+        output = None
+        for shard_index in sorted(set(host_shards)):
+            positions = mx.array(
+                [index for index, value in enumerate(host_shards) if value == shard_index],
+                dtype=mx.int32,
+            )
+            values = self.shards[shard_index](local_ids[positions])
+            if output is None:
+                output = mx.zeros((flat_ids.size, self.dims), dtype=values.dtype)
+            output[positions] = values
+        mx.eval(output)
+        return output.reshape(*shape, self.dims)
+
+
+_SAFETENSORS_NUMPY_DTYPES = {
+    "U32": np.dtype("<u4"),
+    "I32": np.dtype("<i4"),
+    "I64": np.dtype("<i8"),
+    "F16": np.dtype("<f2"),
+    "F32": np.dtype("<f4"),
+    "BF16": np.dtype("<u2"),
+}
+
+
+class _SafeTensorMMap:
+    def __init__(self, path):
+        self.path = Path(path)
+        self._file = self.path.open("rb")
+        raw_header_length = self._file.read(8)
+        if len(raw_header_length) != 8:
+            self.close()
+            raise ValueError(f"Invalid safetensors header in {self.path}")
+        header_length = struct.unpack("<Q", raw_header_length)[0]
+        raw_header = self._file.read(header_length)
+        if len(raw_header) != header_length:
+            self.close()
+            raise ValueError(f"Truncated safetensors header in {self.path}")
+        self._header = json.loads(raw_header)
+        self._data_start = 8 + header_length
+        self._mapping = mmap.mmap(self._file.fileno(), length=0, access=mmap.ACCESS_READ)
+        try:
+            self._mapping.madvise(mmap.MADV_RANDOM)
+        except (AttributeError, OSError):
+            pass
+
+    def tensor_shape(self, key):
+        try:
+            return tuple(self._header[key]["shape"])
+        except KeyError as exc:
+            raise KeyError(f"Tensor {key!r} is missing from {self.path}") from exc
+
+    def rows(self, key, row_indices):
+        try:
+            entry = self._header[key]
+        except KeyError as exc:
+            raise KeyError(f"Tensor {key!r} is missing from {self.path}") from exc
+        dtype_name = entry["dtype"]
+        try:
+            dtype = _SAFETENSORS_NUMPY_DTYPES[dtype_name]
+        except KeyError as exc:
+            raise TypeError(f"Unsupported safetensors dtype {dtype_name!r} for {key}") from exc
+        shape = tuple(entry["shape"])
+        if len(shape) != 2:
+            raise ValueError(f"Sparse PLE tensor {key!r} must be two-dimensional")
+        start, end = entry["data_offsets"]
+        if end - start != math.prod(shape) * dtype.itemsize:
+            raise ValueError(f"Invalid byte range for safetensors tensor {key!r}")
+        view = np.ndarray(
+            shape,
+            dtype=dtype,
+            buffer=self._mapping,
+            offset=self._data_start + start,
+        )
+        copied = np.array(view[np.asarray(row_indices, dtype=np.intp)], copy=True)
+        if dtype_name == "BF16":
+            copied = (copied.astype(np.uint32) << np.uint32(16)).view(np.float32)
+        return copied, dtype_name
+
+    def close(self):
+        mapping = getattr(self, "_mapping", None)
+        if mapping is not None:
+            mapping.close()
+            self._mapping = None
+        file_object = getattr(self, "_file", None)
+        if file_object is not None:
+            file_object.close()
+            self._file = None
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass
+
+
+class DiskBackedShardedEmbedding(nn.Module):
+    def __init__(self, model_path, prefix, num_embeddings, dims, shard_count):
+        super().__init__()
+        if num_embeddings % shard_count:
+            raise ValueError("the padded n-gram vocabulary must divide evenly into shards")
+        self.shard_size = num_embeddings // shard_count
+        self.dims = dims
+        self.shard_count = shard_count
+        self.last_touched_shards = ()
+        self.rows_read = 0
+        self._prefix = prefix
+        self._readers = {}
+        self._tensor_readers = {}
+        self._shard_quantization = []
+
+        model_path = Path(model_path)
+        index_path = model_path / "model.safetensors.index.json"
+        config_path = model_path / "config.json"
+        if not index_path.exists() or not config_path.exists():
+            raise FileNotFoundError("SSD-backed PLE requires config.json and a safetensors index")
+        weight_map = json.loads(index_path.read_text()).get("weight_map", {})
+        quantization = json.loads(config_path.read_text()).get("quantization") or {}
+        default = {
+            "bits": int(quantization.get("bits", 4)),
+            "group_size": int(quantization.get("group_size", 64)),
+            "mode": quantization.get("mode", "affine"),
+        }
+        for shard_index in range(shard_count):
+            module_key = f"{prefix}.shards.{shard_index}"
+            settings = dict(default)
+            settings.update(quantization.get(module_key) or {})
+            bits = int(settings["bits"])
+            group_size = int(settings["group_size"])
+            mode = settings["mode"]
+            if mode != "affine":
+                raise ValueError("SSD-backed PLE supports affine quantization")
+            if dims % group_size or (dims * bits) % 32:
+                raise ValueError(f"Invalid PLE quantization for {module_key}")
+            self._shard_quantization.append((group_size, bits, mode))
+            expected_shapes = {
+                "weight": (self.shard_size, dims * bits // 32),
+                "scales": (self.shard_size, dims // group_size),
+                "biases": (self.shard_size, dims // group_size),
+            }
+            for suffix, expected_shape in expected_shapes.items():
+                key = f"{module_key}.{suffix}"
+                try:
+                    filename = weight_map[key]
+                except KeyError as exc:
+                    raise KeyError(f"SSD-backed PLE tensor {key!r} is absent from the index") from exc
+                reader = self._readers.get(filename)
+                if reader is None:
+                    reader = _SafeTensorMMap(model_path / filename)
+                    self._readers[filename] = reader
+                if reader.tensor_shape(key) != expected_shape:
+                    raise ValueError(
+                        f"Unexpected shape for {key}: {reader.tensor_shape(key)} != {expected_shape}"
+                    )
+                self._tensor_readers[key] = reader
+
+    def _read_rows(self, key, row_indices):
+        array, dtype_name = self._tensor_readers[key].rows(key, row_indices)
+        self.rows_read += len(row_indices)
+        result = mx.array(array)
+        return result.astype(mx.bfloat16) if dtype_name == "BF16" else result
+
+    def __call__(self, ids):
+        shape = ids.shape
+        flat_ids = ids.reshape(-1).astype(mx.int64)
+        mx.eval(flat_ids)
+        host_ids = [int(value) for value in flat_ids.tolist()]
+        if any(value < 0 or value >= self.shard_size * self.shard_count for value in host_ids):
+            raise IndexError("n-gram embedding id is outside the padded vocabulary")
+        touched = tuple(sorted({value // self.shard_size for value in host_ids}))
+        self.last_touched_shards = touched
+        self.rows_read = 0
+        output = None
+        for shard_index in touched:
+            positions_list = [
+                index for index, value in enumerate(host_ids) if value // self.shard_size == shard_index
+            ]
+            local_ids = [host_ids[index] % self.shard_size for index in positions_list]
+            base = f"{self._prefix}.shards.{shard_index}"
+            weight = self._read_rows(f"{base}.weight", local_ids)
+            scales = self._read_rows(f"{base}.scales", local_ids)
+            biases = self._read_rows(f"{base}.biases", local_ids)
+            group_size, bits, mode = self._shard_quantization[shard_index]
+            values = mx.dequantize(
+                weight,
+                scales=scales,
+                biases=biases,
+                group_size=group_size,
+                bits=bits,
+                mode=mode,
+            )
+            if output is None:
+                output = mx.zeros((len(host_ids), self.dims), dtype=values.dtype)
+            output[mx.array(positions_list, dtype=mx.int32)] = values
+        if output is None:
+            output = mx.zeros((0, self.dims), dtype=mx.bfloat16)
+        mx.eval(output)
+        return output.reshape(*shape, self.dims)
+
+
+class Qwen4PLELayer(nn.Module):
+    def __init__(self, args: ModelArgs, ple_index: int, model_layer_index: int):
+        super().__init__()
+        total = args.hc_count * args.hidden_size
+        self.hc_count = args.hc_count
+        self.hidden_size = args.hidden_size
+        self.ple_embedding = Qwen4NGramEmbedding(args, ple_index, model_layer_index)
+        self.key_proj = nn.Linear(args.ple_embed_dim, total, bias=False)
+        self.value_proj = nn.Linear(args.ple_embed_dim, args.hidden_size, bias=False)
+        self.norm_key = Qwen4RMSNorm(total, args.rms_norm_eps, args.hidden_size)
+        self.norm_query = Qwen4RMSNorm(total, args.rms_norm_eps, args.hidden_size)
+        self.norm_conv = Qwen4RMSNorm(total, args.rms_norm_eps, args.hidden_size)
+        self.state_len = (args.ple_conv_kernel_size - 1) * args.ngram_size
+        self.conv1d = nn.Conv1d(
+            total,
+            total,
+            args.ple_conv_kernel_size,
+            dilation=args.ngram_size,
+            groups=total,
+            bias=False,
+        )
+
+    def __call__(self, x, input_ids, cache=None, mask=None):
+        embeddings = self.ple_embedding(input_ids, cache)
+        key = self.norm_key(self.key_proj(embeddings)).reshape(*x.shape[:-1], self.hc_count, self.hidden_size)
+        value = self.value_proj(embeddings)
+        query = self.norm_query(x).reshape(*x.shape[:-1], self.hc_count, self.hidden_size)
+        gate = mx.sum(key * query, axis=-1, keepdims=True) / math.sqrt(self.hidden_size)
+        gate = mx.sign(gate) * mx.sqrt(mx.maximum(mx.abs(gate), 1e-6))
+        gated = (mx.sigmoid(gate) * value[..., None, :]).reshape(*x.shape)
+        normalized = self.norm_conv(gated)
+        if mask is not None:
+            gated = mx.where(mask[..., None], gated, 0)
+            normalized = mx.where(mask[..., None], normalized, 0)
+        if cache is not None and cache[2] is not None:
+            state = cache[2]
+        else:
+            state = mx.zeros((x.shape[0], self.state_len, x.shape[-1]), dtype=x.dtype)
+        conv_input = mx.concatenate([state, normalized], axis=1)
+        if cache is not None:
+            cache[2] = mx.contiguous(conv_input[:, -self.state_len :, :])
+        return gated + nn.silu(self.conv1d(conv_input))
+
+
+class Qwen4DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, index: int):
+        super().__init__()
+        self.is_linear = args.layer_types[index] == "linear_attention"
+        if self.is_linear:
+            self.linear_attn = Qwen4GatedDeltaNet(args)
+        else:
+            self.self_attn = Qwen4Attention(args)
+        self.mlp = Qwen4SparseMoeBlock(args)
+        one_indexed = index + 1
+        self.ple = Qwen4PLELayer(args, args.ple_layer_ids.index(one_indexed), index) if one_indexed in args.ple_layer_ids else None
+        self.attn_hyper_connection = Qwen4GatedResidual(args)
+        self.mlp_hyper_connection = Qwen4GatedResidual(args)
+
+    def __call__(self, x, input_ids, mask=None, cache=None, n_confirmed=0):
+        if self.ple is not None:
+            x = x + self.ple(x, input_ids, cache, mask)
+        mixed, residual, injection = self.attn_hyper_connection(x)
+        output = (
+            self.linear_attn(mixed, mask, cache, n_confirmed=n_confirmed)
+            if self.is_linear
+            else self.self_attn(mixed, cache)
+        )
+        x = residual + (output[..., None, :] * injection[..., None]).reshape(*residual.shape)
+        mixed, residual, injection = self.mlp_hyper_connection(x)
+        output = self.mlp(mixed)
+        return residual + (output[..., None, :] * injection[..., None]).reshape(*residual.shape)
+
+
+class Qwen4MTPModule(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.hc_count = args.hc_count
+        total = self.hc_count * self.hidden_size
+        self.pre_fc_norm_embedding = Qwen4RMSNorm(
+            self.hidden_size,
+            args.rms_norm_eps,
+        )
+        self.pre_fc_norm_hidden = Qwen4RMSNorm(total, args.rms_norm_eps)
+        self.fc_embedding = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.fc_hidden = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        mtp_args = copy(args)
+        mtp_args.num_hidden_layers = 1
+        mtp_args.layer_types = ["full_attention"]
+        mtp_args.ple_layer_ids = []
+        self.layers = [Qwen4DecoderLayer(mtp_args, 0)]
+        self.hyper_connection_mixer = Qwen4GatedResidual(mtp_args, combine=False)
+
+    def fuse_inputs(self, input_embeds, hidden_states):
+        input_embeds = self.fc_embedding(
+            self.pre_fc_norm_embedding(input_embeds)
+        )
+        original_shape = hidden_states.shape
+        streams = self.pre_fc_norm_hidden(hidden_states).reshape(
+            *hidden_states.shape[:-1],
+            self.hc_count,
+            self.hidden_size,
+        )
+        streams = self.fc_hidden(streams)
+        return (streams + input_embeds[..., None, :]).reshape(original_shape)
+
+    def __call__(self, hidden_states, next_token_ids, embed_tokens, cache=None):
+        hidden_states = self.fuse_inputs(
+            embed_tokens(next_token_ids),
+            hidden_states,
+        )
+        if cache is None:
+            cache = [None] * len(self.layers)
+        for layer, layer_cache in zip(self.layers, cache):
+            hidden_states = layer(
+                hidden_states,
+                next_token_ids,
+                cache=layer_cache,
+            )
+        return self.hyper_connection_mixer(hidden_states), hidden_states
+
+
+class Qwen4TextModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [Qwen4DecoderLayer(args, index) for index in range(args.num_hidden_layers)]
+        self.hyper_connection_mixer = Qwen4GatedResidual(args, combine=False)
+
+    def __call__(self, inputs, cache=None, return_hidden=False, n_confirmed=0):
+        hidden = self.embed_tokens(inputs)
+        hidden = mx.tile(hidden, (1, 1, self.args.hc_count))
+        if cache is None:
+            cache = [None] * len(self.layers)
+        linear_index = next((i for i, layer in enumerate(self.layers) if layer.is_linear), None)
+        mask = create_ssm_mask(hidden, cache[linear_index]) if linear_index is not None else None
+        for layer, layer_cache in zip(self.layers, cache):
+            hidden = layer(
+                hidden,
+                inputs,
+                mask,
+                layer_cache,
+                n_confirmed=n_confirmed,
+            )
+        output = self.hyper_connection_mixer(hidden)
+        if return_hidden:
+            return output, hidden
+        return output
+
+
+def _decode_block_fp8(weight, scale, block=128):
+    rows, columns = weight.shape
+    grid_rows, grid_columns = scale.shape
+    padded_rows, padded_columns = grid_rows * block, grid_columns * block
+    decoded = mx.from_fp8(weight, dtype=mx.float32)
+    decoded = mx.pad(decoded, ((0, padded_rows - rows), (0, padded_columns - columns)))
+    decoded = decoded.reshape(grid_rows, block, grid_columns, block)
+    decoded = decoded * scale.astype(mx.float32)[:, None, :, None]
+    return decoded.reshape(padded_rows, padded_columns)[:rows, :columns].astype(mx.bfloat16)
+
+
+def register_oq_virtual_tensors(index, config):
+    if config.get("model_type") != "qwen4_exp":
+        return 0
+    text_config = config.get("text_config", config)
+    registrations = 0
+    shard_count = int(text_config.get("split_ngram_parts", 0))
+    for one_indexed_layer in text_config.get("ple_layer_ids", []):
+        prefix = (
+            f"model.language_model.layers.{one_indexed_layer - 1}.ple."
+            "ple_embedding.ngram_embedding"
+        )
+        scale_key = prefix + ".weight_scale"
+        if index.source_shape(scale_key) is None:
+            continue
+        for shard_index in range(shard_count):
+            weight_key = f"{prefix}.shard_{shard_index}.weight"
+            shape = index.source_shape(weight_key)
+            if shape is None:
+                raise ValueError(f"Missing Qwen4 PLE shard: {weight_key}")
+
+            def materialize(weight_key=weight_key, scale_key=scale_key):
+                weight = index.load_source(weight_key)
+                scale = index.load_source(scale_key)
+                value = mx.from_fp8(weight, dtype=mx.bfloat16) * scale
+                mx.eval(value)
+                return value
+
+            index.register_virtual(
+                weight_key,
+                shape,
+                "BF16",
+                materialize,
+                hides=(weight_key, scale_key),
+            )
+            registrations += 1
+    return registrations
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = Qwen4TextModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        mtp_active = False
+        if args.mtp_num_hidden_layers:
+            try:
+                from omlx.patches.mlx_lm_mtp import is_mtp_active
+
+                mtp_active = is_mtp_active()
+            except ImportError:
+                mtp_active = False
+        self._omlx_mtp_decode_enabled = bool(mtp_active)
+        if mtp_active:
+            if args.mtp_num_hidden_layers != 1:
+                raise ValueError("Qwen4 requires exactly one MTP decoder layer")
+            self.mtp = Qwen4MTPModule(args)
+            self._omlx_mtp_chain = True
+            self._omlx_mtp_depth = 1
+            self._omlx_mtp_head_prenorm = True
+
+    def __call__(self, inputs, cache=None, return_hidden=False, n_confirmed=0):
+        need_hidden = return_hidden or hasattr(self, "mtp")
+        result = self.model(
+            inputs,
+            cache,
+            return_hidden=need_hidden,
+            n_confirmed=n_confirmed,
+        )
+        if need_hidden:
+            output, hidden = result
+        else:
+            output = result
+        if self.args.tie_word_embeddings:
+            logits = self.model.embed_tokens.as_linear(output)
+        else:
+            logits = self.lm_head(output)
+        if (
+            hasattr(self, "mtp")
+            and not return_hidden
+            and not n_confirmed
+        ):
+            try:
+                from omlx.patches.mlx_lm_mtp import prompt_priming
+
+                prompt_priming.maybe_capture(self, inputs, hidden, cache)
+            except Exception:
+                logger.debug("Qwen4 MTP prompt priming failed", exc_info=True)
+        if return_hidden:
+            return logits, hidden
+        return logits
+
+    def mtp_forward(
+        self,
+        hidden_states,
+        next_token_ids,
+        mtp_cache,
+        return_hidden=False,
+        logits_keep=0,
+    ):
+        if not hasattr(self, "mtp"):
+            raise RuntimeError("Qwen4 MTP forward called without an attached head")
+        output, hidden = self.mtp(
+            hidden_states,
+            next_token_ids,
+            self.model.embed_tokens,
+            mtp_cache,
+        )
+        if logits_keep and output.shape[1] > logits_keep:
+            output = output[:, -logits_keep:, :]
+        if self.args.tie_word_embeddings:
+            logits = self.model.embed_tokens.as_linear(output)
+        else:
+            logits = self.lm_head(output)
+        if return_hidden:
+            return logits, hidden
+        return logits
+
+    def make_mtp_cache(self):
+        if not hasattr(self, "mtp"):
+            return []
+        return [QSAKVCache() for _ in self.mtp.layers]
+
+    def mtp_partial_rollback(self, cache, accepted, num_drafts):
+        if len(cache) != len(self.layers):
+            return False
+        trim = int(num_drafts) - int(accepted)
+        if trim <= 0:
+            return True
+        keep = 1 + int(accepted)
+        for layer, layer_cache in zip(self.layers, cache):
+            if layer.is_linear:
+                if getattr(layer_cache, "rollback_state", None) is None:
+                    return False
+                if getattr(layer_cache, "_mtp_draft_stash", None) is None:
+                    return False
+            elif not layer_cache.is_trimmable():
+                return False
+        for layer, layer_cache in zip(self.layers, cache):
+            if layer.is_linear:
+                conv_state, state = layer_cache.rollback_state
+                mixed, a, b, mask = layer_cache._mtp_draft_stash
+                if mask is not None:
+                    mask = mask[:, :keep]
+                _, next_conv_state, next_state = layer.linear_attn._process_chunk(
+                    mixed[:, :keep],
+                    a[:, :keep],
+                    b[:, :keep],
+                    conv_state,
+                    state,
+                    mask,
+                )
+                layer_cache[0] = next_conv_state
+                layer_cache[1] = next_state
+                layer_cache.rollback_state = None
+                layer_cache._mtp_draft_stash = None
+            else:
+                layer_cache.trim(trim)
+        return True
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [ArraysCache(size=4) if layer.is_linear else QSAKVCache() for layer in self.layers]
+
+    def sanitize(self, weights):
+        cleaned = {}
+        for key, value in weights.items():
+            if key.startswith("mtp.") and not hasattr(self, "mtp"):
+                continue
+            if key.startswith("model.visual."):
+                continue
+            if key.startswith("model.language_model."):
+                key = "model." + key[len("model.language_model.") :]
+            cleaned[key] = value
+        weights = cleaned
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+        expert_prefixes = [
+            f"model.layers.{layer_index}.mlp"
+            for layer_index in range(self.args.num_hidden_layers)
+        ]
+        if hasattr(self, "mtp"):
+            expert_prefixes.extend(
+                f"mtp.layers.{layer_index}.mlp"
+                for layer_index in range(self.args.mtp_num_hidden_layers)
+            )
+        for prefix in expert_prefixes:
+            if f"{prefix}.experts.0.up_proj.weight" not in weights:
+                continue
+            for projection in ("up_proj", "gate_proj", "down_proj"):
+                values = []
+                for expert in range(self.args.num_experts):
+                    key = f"{prefix}.experts.{expert}.{projection}.weight"
+                    scale_key = key + "_scale_inv"
+                    value = weights.pop(key)
+                    if scale_key in weights:
+                        value = _decode_block_fp8(value, weights.pop(scale_key))
+                    values.append(value)
+                weights[f"{prefix}.switch_mlp.{projection}.weight"] = mx.stack(values)
+        for layer_index in range(self.args.num_hidden_layers):
+            base = f"model.layers.{layer_index}.ple.ple_embedding.ngram_embedding"
+            layer = self.model.layers[layer_index]
+            disk_backed = (
+                layer.ple is not None
+                and isinstance(
+                    layer.ple.ple_embedding.ngram_embedding,
+                    DiskBackedShardedEmbedding,
+                )
+            )
+            if disk_backed:
+                for key in list(weights):
+                    if key.startswith(base + ".shards."):
+                        weights.pop(key)
+            shard_keys = [key for key in weights if key.startswith(base + ".shard_") and key.endswith(".weight")]
+            if not shard_keys:
+                continue
+            shard_keys.sort(key=lambda key: int(key.split(".shard_")[1].split(".")[0]))
+            scale = weights.pop(base + ".weight_scale", None)
+            for shard_index, key in enumerate(shard_keys):
+                value = weights.pop(key)
+                if scale is not None:
+                    value = mx.from_fp8(value, dtype=mx.bfloat16) * scale
+                weights[f"{base}.shards.{shard_index}.weight"] = value
+        for key in list(weights):
+            value = weights[key]
+            if key.endswith("conv1d.weight") and value.ndim == 3 and value.shape[-1] != 1:
+                weights[key] = value.moveaxis(2, 1)
+        return weights
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, module):
+            if "ple_embedding.ngram_embedding.shards." in path:
+                return {"group_size": 32, "bits": 4}
+            if path.endswith("mlp.gate") or path.endswith("shared_expert_gate"):
+                return {"group_size": 64, "bits": 8}
+            return hasattr(module, "to_quantized")
+
+        return predicate

omlx_support/qwen4_mtp_integration.py

@@ -0,0 +1,18 @@
+def register_qwen4_mtp_integration():
+    try:
+        from omlx.utils import model_loading
+    except ImportError:
+        return False
+    if getattr(model_loading, "_qwen4_mtp_compatible_installed", False):
+        return True
+
+    original = model_loading._is_mtp_compatible
+
+    def compatible(config, model_type):
+        if model_type == "qwen4_exp":
+            return model_loading._has_mtp_heads(config)
+        return original(config, model_type)
+
+    model_loading._is_mtp_compatible = compatible
+    model_loading._qwen4_mtp_compatible_installed = True
+    return True

omlx_support/serve

@@ -0,0 +1,33 @@
+#!/bin/zsh
+set -euo pipefail
+
+support_root=${0:A:h}
+model_root=${support_root:h}
+models_root=${model_root:h}
+base_path=${OMLX_BASE_PATH:-$model_root/.omlx}
+cache_dir=${OMLX_CACHE_DIR:-$base_path/cache}
+api_key=${OMLX_API_KEY:-omlx}
+port=${OMLX_PORT:-8766}
+
+mkdir -p "$cache_dir"
+if [[ ! -f "$base_path/model_settings.json" ]]; then
+  cp "$support_root/model_settings.json" "$base_path/model_settings.json"
+fi
+export OMLX_QWEN4_PLE_MODE=mmap
+export OMLX_QWEN4_PLE_MODEL_PATH="$model_root"
+export OMLX_MTP_PROMPT_PRIMING=1
+export OMLX_MTP_PRIME_WINDOW=8192
+
+exec /bin/zsh "$support_root/with-omlx-python" -m omlx.cli serve \
+  --model-dir "$models_root" \
+  --host 127.0.0.1 \
+  --port "$port" \
+  --max-concurrent-requests 1 \
+  --memory-guard balanced \
+  --paged-ssd-cache-dir "$cache_dir" \
+  --paged-ssd-cache-max-size 128GB \
+  --hot-cache-max-size 0 \
+  --initial-cache-blocks 1 \
+  --no-hf-cache \
+  --base-path "$base_path" \
+  --api-key "$api_key"

omlx_support/sitecustomize.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+
+import mlx_lm.models
+
+
+model_path = str(Path(__file__).resolve().parent)
+if model_path not in mlx_lm.models.__path__:
+    mlx_lm.models.__path__.insert(0, model_path)
+
+from qwen4_cache_integration import register_qwen4_cache_integration
+from qwen4_mtp_integration import register_qwen4_mtp_integration
+
+
+register_qwen4_cache_integration()
+register_qwen4_mtp_integration()

omlx_support/with-omlx-python

@@ -0,0 +1,22 @@
+#!/bin/zsh
+set -euo pipefail
+
+support_root=${0:A:h}
+model_root=${support_root:h}
+app_resources=${OMLX_APP_RESOURCES:-/Applications/oMLX.app/Contents/Resources}
+python_root=$app_resources/Python
+cpython_root=$python_root/cpython-3.11
+mlx_site=$python_root/framework-mlx-base/lib/python3.11/site-packages
+mlx_runtime=${OMLX_MLX_RUNTIME:-$model_root/.mlx-runtime}
+
+if [[ ! -d "$mlx_runtime/mlx" ]]; then
+  print -u2 "MLX 0.32.1 is missing from $mlx_runtime"
+  print -u2 "Install it with: uv pip install --target $mlx_runtime mlx==0.32.1 mlx-metal==0.32.1"
+  exit 1
+fi
+
+export PYTHONHOME=$cpython_root
+export PYTHONDONTWRITEBYTECODE=1
+export PYTHONPATH="$support_root:$mlx_runtime:$app_resources:$mlx_site"
+
+exec "$cpython_root/bin/python3" "$@"

oq_imatrix_report.json

@@ -0,0 +1,870 @@
+{
+  "enabled": true,
+  "cache_path": null,
+  "cache_reused": true,
+  "entry_count": 914,
+  "calib_dataset": "oqe_code_multilingual",
+  "collection": {
+    "dataset": "oqe_code_multilingual",
+    "requested_samples": 128,
+    "seq_length": 512,
+    "adaptive": true,
+    "adaptive_step_samples": 128,
+    "adaptive_max_samples": 1024,
+    "available_samples": 1024,
+    "micro_batch_size": 6,
+    "micro_batches": 176,
+    "batch_plan": {
+      "micro_batch_size": 6,
+      "estimated_sample_bytes": 52428800,
+      "capture_budget_bytes": 362029891,
+      "system_available_bytes": 103756529664,
+      "metal_available_bytes": 115443233112,
+      "live_available_bytes": 103756529664,
+      "model_bytes": 67553540501,
+      "remaining_available_bytes": 36202989163,
+      "fits_one_sample": true,
+      "hidden_size": 2560,
+      "num_experts": 512,
+      "top_k": 10,
+      "gemma4_state_bytes": 0,
+      "num_hidden_layers": 48,
+      "num_kv_shared_layers": 0,
+      "per_layer_input_size": 0
+    },
+    "processed_samples": 1024,
+    "installed_modules": 917,
+    "capture_module_classes": {
+      "Linear": 388,
+      "QuantizedLinear": 385,
+      "QuantizedSwitchLinear": 144
+    },
+    "switch_capture_modules": 144,
+    "requires_expert_counts": true,
+    "coverage_sufficient": false,
+    "collection_sufficient": false,
+    "coverage": {
+      "has_expert_counts": true,
+      "expert_modules": 144,
+      "total_experts": 73728,
+      "active_experts": 73665,
+      "zero_count_experts": 63,
+      "active_ratio": 0.9991455078125,
+      "min_count": 0,
+      "p05_count": 717.0,
+      "p10_count": 1432.0,
+      "median_count": 7353.0,
+      "max_count": 341355,
+      "min_required_count": 16,
+      "required_percentile": 5
+    },
+    "rounds": [
+      {
+        "processed_samples": 128,
+        "coverage_sufficient": false,
+        "collection_sufficient": false,
+        "coverage": {
+          "has_expert_counts": true,
+          "expert_modules": 144,
+          "total_experts": 73728,
+          "active_experts": 73530,
+          "zero_count_experts": 198,
+          "active_ratio": 0.997314453125,
+          "min_count": 0,
+          "p05_count": 71.0,
+          "p10_count": 158.0,
+          "median_count": 913.0,
+          "max_count": 41723,
+          "min_required_count": 16,
+          "required_percentile": 5
+        }
+      },
+      {
+        "processed_samples": 256,
+        "coverage_sufficient": false,
+        "collection_sufficient": false,
+        "coverage": {
+          "has_expert_counts": true,
+          "expert_modules": 144,
+          "total_experts": 73728,
+          "active_experts": 73587,
+          "zero_count_experts": 141,
+          "active_ratio": 0.9980875651041666,
+          "min_count": 0,
+          "p05_count": 174.0,
+          "p10_count": 353.0,
+          "median_count": 1830.0,
+          "max_count": 83264,
+          "min_required_count": 16,
+          "required_percentile": 5
+        }
+      },
+      {
+        "processed_samples": 384,
+        "coverage_sufficient": false,
+        "collection_sufficient": false,
+        "coverage": {
+          "has_expert_counts": true,
+          "expert_modules": 144,
+          "total_experts": 73728,
+          "active_experts": 73629,
+          "zero_count_experts": 99,
+          "active_ratio": 0.9986572265625,
+          "min_count": 0,
+          "p05_count": 285.35000000000036,
+          "p10_count": 553.0,
+          "median_count": 2755.0,
+          "max_count": 124541,
+          "min_required_count": 16,
+          "required_percentile": 5
+        }
+      },
+      {
+        "processed_samples": 512,
+        "coverage_sufficient": false,
+        "collection_sufficient": false,
+        "coverage": {
+          "has_expert_counts": true,
+          "expert_modules": 144,
+          "total_experts": 73728,
+          "active_experts": 73641,
+          "zero_count_experts": 87,
+          "active_ratio": 0.9988199869791666,
+          "min_count": 0,
+          "p05_count": 369.35000000000036,
+          "p10_count": 737.0,
+          "median_count": 3687.5,
+          "max_count": 168254,
+          "min_required_count": 16,
+          "required_percentile": 5
+        }
+      },
+      {
+        "processed_samples": 640,
+        "coverage_sufficient": false,
+        "collection_sufficient": false,
+        "coverage": {
+          "has_expert_counts": true,
+          "expert_modules": 144,
+          "total_experts": 73728,
+          "active_experts": 73653,
+          "zero_count_experts": 75,
+          "active_ratio": 0.9989827473958334,
+          "min_count": 0,
+          "p05_count": 453.35000000000036,
+          "p10_count": 904.0,
+          "median_count": 4619.0,
+          "max_count": 213373,
+          "min_required_count": 16,
+          "required_percentile": 5
+        }
+      },
+      {
+        "processed_samples": 768,
+        "coverage_sufficient": false,
+        "collection_sufficient": false,
+        "coverage": {
+          "has_expert_counts": true,
+          "expert_modules": 144,
+          "total_experts": 73728,
+          "active_experts": 73656,
+          "zero_count_experts": 72,
+          "active_ratio": 0.9990234375,
+          "min_count": 0,
+          "p05_count": 539.0,
+          "p10_count": 1076.0,
+          "median_count": 5501.5,
+          "max_count": 255296,
+          "min_required_count": 16,
+          "required_percentile": 5
+        }
+      },
+      {
+        "processed_samples": 896,
+        "coverage_sufficient": false,
+        "collection_sufficient": false,
+        "coverage": {
+          "has_expert_counts": true,
+          "expert_modules": 144,
+          "total_experts": 73728,
+          "active_experts": 73659,
+          "zero_count_experts": 69,
+          "active_ratio": 0.9990641276041666,
+          "min_count": 0,
+          "p05_count": 628.0,
+          "p10_count": 1257.0,
+          "median_count": 6420.0,
+          "max_count": 298656,
+          "min_required_count": 16,
+          "required_percentile": 5
+        }
+      },
+      {
+        "processed_samples": 1024,
+        "coverage_sufficient": false,
+        "collection_sufficient": false,
+        "coverage": {
+          "has_expert_counts": true,
+          "expert_modules": 144,
+          "total_experts": 73728,
+          "active_experts": 73665,
+          "zero_count_experts": 63,
+          "active_ratio": 0.9991455078125,
+          "min_count": 0,
+          "p05_count": 717.0,
+          "p10_count": 1432.0,
+          "median_count": 7353.0,
+          "max_count": 341355,
+          "min_required_count": 16,
+          "required_percentile": 5
+        }
+      }
+    ]
+  },
+  "expert_coverage": {
+    "has_expert_counts": true,
+    "expert_modules": 144,
+    "total_experts": 73728,
+    "active_experts": 73665,
+    "zero_count_experts": 63,
+    "active_ratio": 0.9991455078125,
+    "min_count": 0,
+    "p05_count": 717.0,
+    "p10_count": 1432.0,
+    "median_count": 7353.0,
+    "max_count": 341355,
+    "min_required_count": 16,
+    "required_percentile": 5
+  },
+  "applied": [
+    "model.layers.0.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.0.linear_attn.in_proj_a",
+    "model.layers.0.linear_attn.in_proj_b",
+    "model.layers.0.linear_attn.in_proj_qkv",
+    "model.layers.0.linear_attn.in_proj_z",
+    "model.layers.0.linear_attn.out_proj",
+    "model.layers.0.mlp.shared_expert.down_proj",
+    "model.layers.0.mlp.shared_expert.gate_proj",
+    "model.layers.0.mlp.shared_expert.up_proj",
+    "model.layers.0.mlp.switch_mlp.down_proj",
+    "model.layers.0.mlp.switch_mlp.gate_proj",
+    "model.layers.0.mlp.switch_mlp.up_proj",
+    "model.layers.0.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.1.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.1.linear_attn.in_proj_a",
+    "model.layers.1.linear_attn.in_proj_b",
+    "model.layers.1.linear_attn.in_proj_qkv",
+    "model.layers.1.linear_attn.in_proj_z",
+    "model.layers.1.linear_attn.out_proj",
+    "model.layers.1.mlp.shared_expert.down_proj",
+    "model.layers.1.mlp.shared_expert.gate_proj",
+    "model.layers.1.mlp.shared_expert.up_proj",
+    "model.layers.1.mlp.switch_mlp.down_proj",
+    "model.layers.1.mlp.switch_mlp.gate_proj",
+    "model.layers.1.mlp.switch_mlp.up_proj",
+    "model.layers.1.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.1.ple.key_proj",
+    "model.layers.1.ple.value_proj",
+    "model.layers.10.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.10.linear_attn.in_proj_a",
+    "model.layers.10.linear_attn.in_proj_b",
+    "model.layers.10.linear_attn.in_proj_qkv",
+    "model.layers.10.linear_attn.in_proj_z",
+    "model.layers.10.linear_attn.out_proj",
+    "model.layers.10.mlp.shared_expert.down_proj",
+    "model.layers.10.mlp.shared_expert.gate_proj",
+    "model.layers.10.mlp.shared_expert.up_proj",
+    "model.layers.10.mlp.switch_mlp.down_proj",
+    "model.layers.10.mlp.switch_mlp.gate_proj",
+    "model.layers.10.mlp.switch_mlp.up_proj",
+    "model.layers.10.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.11.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.11.mlp.shared_expert.down_proj",
+    "model.layers.11.mlp.shared_expert.gate_proj",
+    "model.layers.11.mlp.shared_expert.up_proj",
+    "model.layers.11.mlp.switch_mlp.down_proj",
+    "model.layers.11.mlp.switch_mlp.gate_proj",
+    "model.layers.11.mlp.switch_mlp.up_proj",
+    "model.layers.11.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.11.self_attn.indexer.index_qk_proj",
+    "model.layers.11.self_attn.k_proj",
+    "model.layers.11.self_attn.o_proj",
+    "model.layers.11.self_attn.q_proj",
+    "model.layers.11.self_attn.v_proj",
+    "model.layers.12.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.12.linear_attn.in_proj_a",
+    "model.layers.12.linear_attn.in_proj_b",
+    "model.layers.12.linear_attn.in_proj_qkv",
+    "model.layers.12.linear_attn.in_proj_z",
+    "model.layers.12.linear_attn.out_proj",
+    "model.layers.12.mlp.shared_expert.down_proj",
+    "model.layers.12.mlp.shared_expert.gate_proj",
+    "model.layers.12.mlp.shared_expert.up_proj",
+    "model.layers.12.mlp.switch_mlp.down_proj",
+    "model.layers.12.mlp.switch_mlp.gate_proj",
+    "model.layers.12.mlp.switch_mlp.up_proj",
+    "model.layers.12.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.13.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.13.linear_attn.in_proj_a",
+    "model.layers.13.linear_attn.in_proj_b",
+    "model.layers.13.linear_attn.in_proj_qkv",
+    "model.layers.13.linear_attn.in_proj_z",
+    "model.layers.13.linear_attn.out_proj",
+    "model.layers.13.mlp.shared_expert.down_proj",
+    "model.layers.13.mlp.shared_expert.gate_proj",
+    "model.layers.13.mlp.shared_expert.up_proj",
+    "model.layers.13.mlp.switch_mlp.down_proj",
+    "model.layers.13.mlp.switch_mlp.gate_proj",
+    "model.layers.13.mlp.switch_mlp.up_proj",
+    "model.layers.13.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.14.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.14.linear_attn.in_proj_a",
+    "model.layers.14.linear_attn.in_proj_b",
+    "model.layers.14.linear_attn.in_proj_qkv",
+    "model.layers.14.linear_attn.in_proj_z",
+    "model.layers.14.linear_attn.out_proj",
+    "model.layers.14.mlp.shared_expert.down_proj",
+    "model.layers.14.mlp.shared_expert.gate_proj",
+    "model.layers.14.mlp.shared_expert.up_proj",
+    "model.layers.14.mlp.switch_mlp.down_proj",
+    "model.layers.14.mlp.switch_mlp.gate_proj",
+    "model.layers.14.mlp.switch_mlp.up_proj",
+    "model.layers.14.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.15.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.15.mlp.shared_expert.down_proj",
+    "model.layers.15.mlp.shared_expert.gate_proj",
+    "model.layers.15.mlp.shared_expert.up_proj",
+    "model.layers.15.mlp.switch_mlp.down_proj",
+    "model.layers.15.mlp.switch_mlp.gate_proj",
+    "model.layers.15.mlp.switch_mlp.up_proj",
+    "model.layers.15.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.15.self_attn.indexer.index_qk_proj",
+    "model.layers.15.self_attn.k_proj",
+    "model.layers.15.self_attn.o_proj",
+    "model.layers.15.self_attn.q_proj",
+    "model.layers.15.self_attn.v_proj",
+    "model.layers.16.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.16.linear_attn.in_proj_a",
+    "model.layers.16.linear_attn.in_proj_b",
+    "model.layers.16.linear_attn.in_proj_qkv",
+    "model.layers.16.linear_attn.in_proj_z",
+    "model.layers.16.linear_attn.out_proj",
+    "model.layers.16.mlp.shared_expert.down_proj",
+    "model.layers.16.mlp.shared_expert.gate_proj",
+    "model.layers.16.mlp.shared_expert.up_proj",
+    "model.layers.16.mlp.switch_mlp.down_proj",
+    "model.layers.16.mlp.switch_mlp.gate_proj",
+    "model.layers.16.mlp.switch_mlp.up_proj",
+    "model.layers.16.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.17.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.17.linear_attn.in_proj_a",
+    "model.layers.17.linear_attn.in_proj_b",
+    "model.layers.17.linear_attn.in_proj_qkv",
+    "model.layers.17.linear_attn.in_proj_z",
+    "model.layers.17.linear_attn.out_proj",
+    "model.layers.17.mlp.shared_expert.down_proj",
+    "model.layers.17.mlp.shared_expert.gate_proj",
+    "model.layers.17.mlp.shared_expert.up_proj",
+    "model.layers.17.mlp.switch_mlp.down_proj",
+    "model.layers.17.mlp.switch_mlp.gate_proj",
+    "model.layers.17.mlp.switch_mlp.up_proj",
+    "model.layers.17.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.18.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.18.linear_attn.in_proj_a",
+    "model.layers.18.linear_attn.in_proj_b",
+    "model.layers.18.linear_attn.in_proj_qkv",
+    "model.layers.18.linear_attn.in_proj_z",
+    "model.layers.18.linear_attn.out_proj",
+    "model.layers.18.mlp.shared_expert.down_proj",
+    "model.layers.18.mlp.shared_expert.gate_proj",
+    "model.layers.18.mlp.shared_expert.up_proj",
+    "model.layers.18.mlp.switch_mlp.down_proj",
+    "model.layers.18.mlp.switch_mlp.gate_proj",
+    "model.layers.18.mlp.switch_mlp.up_proj",
+    "model.layers.18.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.19.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.19.mlp.shared_expert.down_proj",
+    "model.layers.19.mlp.shared_expert.gate_proj",
+    "model.layers.19.mlp.shared_expert.up_proj",
+    "model.layers.19.mlp.switch_mlp.down_proj",
+    "model.layers.19.mlp.switch_mlp.gate_proj",
+    "model.layers.19.mlp.switch_mlp.up_proj",
+    "model.layers.19.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.19.self_attn.indexer.index_qk_proj",
+    "model.layers.19.self_attn.k_proj",
+    "model.layers.19.self_attn.o_proj",
+    "model.layers.19.self_attn.q_proj",
+    "model.layers.19.self_attn.v_proj",
+    "model.layers.2.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.2.linear_attn.in_proj_a",
+    "model.layers.2.linear_attn.in_proj_b",
+    "model.layers.2.linear_attn.in_proj_qkv",
+    "model.layers.2.linear_attn.in_proj_z",
+    "model.layers.2.linear_attn.out_proj",
+    "model.layers.2.mlp.shared_expert.down_proj",
+    "model.layers.2.mlp.shared_expert.gate_proj",
+    "model.layers.2.mlp.shared_expert.up_proj",
+    "model.layers.2.mlp.switch_mlp.down_proj",
+    "model.layers.2.mlp.switch_mlp.gate_proj",
+    "model.layers.2.mlp.switch_mlp.up_proj",
+    "model.layers.2.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.20.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.20.linear_attn.in_proj_a",
+    "model.layers.20.linear_attn.in_proj_b",
+    "model.layers.20.linear_attn.in_proj_qkv",
+    "model.layers.20.linear_attn.in_proj_z",
+    "model.layers.20.linear_attn.out_proj",
+    "model.layers.20.mlp.shared_expert.down_proj",
+    "model.layers.20.mlp.shared_expert.gate_proj",
+    "model.layers.20.mlp.shared_expert.up_proj",
+    "model.layers.20.mlp.switch_mlp.down_proj",
+    "model.layers.20.mlp.switch_mlp.gate_proj",
+    "model.layers.20.mlp.switch_mlp.up_proj",
+    "model.layers.20.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.21.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.21.linear_attn.in_proj_a",
+    "model.layers.21.linear_attn.in_proj_b",
+    "model.layers.21.linear_attn.in_proj_qkv",
+    "model.layers.21.linear_attn.in_proj_z",
+    "model.layers.21.linear_attn.out_proj",
+    "model.layers.21.mlp.shared_expert.down_proj",
+    "model.layers.21.mlp.shared_expert.gate_proj",
+    "model.layers.21.mlp.shared_expert.up_proj",
+    "model.layers.21.mlp.switch_mlp.down_proj",
+    "model.layers.21.mlp.switch_mlp.gate_proj",
+    "model.layers.21.mlp.switch_mlp.up_proj",
+    "model.layers.21.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.22.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.22.linear_attn.in_proj_a",
+    "model.layers.22.linear_attn.in_proj_b",
+    "model.layers.22.linear_attn.in_proj_qkv",
+    "model.layers.22.linear_attn.in_proj_z",
+    "model.layers.22.linear_attn.out_proj",
+    "model.layers.22.mlp.shared_expert.down_proj",
+    "model.layers.22.mlp.shared_expert.gate_proj",
+    "model.layers.22.mlp.shared_expert.up_proj",
+    "model.layers.22.mlp.switch_mlp.down_proj",
+    "model.layers.22.mlp.switch_mlp.gate_proj",
+    "model.layers.22.mlp.switch_mlp.up_proj",
+    "model.layers.22.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.23.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.23.mlp.shared_expert.down_proj",
+    "model.layers.23.mlp.shared_expert.gate_proj",
+    "model.layers.23.mlp.shared_expert.up_proj",
+    "model.layers.23.mlp.switch_mlp.down_proj",
+    "model.layers.23.mlp.switch_mlp.gate_proj",
+    "model.layers.23.mlp.switch_mlp.up_proj",
+    "model.layers.23.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.23.self_attn.indexer.index_qk_proj",
+    "model.layers.23.self_attn.k_proj",
+    "model.layers.23.self_attn.o_proj",
+    "model.layers.23.self_attn.q_proj",
+    "model.layers.23.self_attn.v_proj",
+    "model.layers.24.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.24.linear_attn.in_proj_a",
+    "model.layers.24.linear_attn.in_proj_b",
+    "model.layers.24.linear_attn.in_proj_qkv",
+    "model.layers.24.linear_attn.in_proj_z",
+    "model.layers.24.linear_attn.out_proj",
+    "model.layers.24.mlp.shared_expert.down_proj",
+    "model.layers.24.mlp.shared_expert.gate_proj",
+    "model.layers.24.mlp.shared_expert.up_proj",
+    "model.layers.24.mlp.switch_mlp.down_proj",
+    "model.layers.24.mlp.switch_mlp.gate_proj",
+    "model.layers.24.mlp.switch_mlp.up_proj",
+    "model.layers.24.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.25.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.25.linear_attn.in_proj_a",
+    "model.layers.25.linear_attn.in_proj_b",
+    "model.layers.25.linear_attn.in_proj_qkv",
+    "model.layers.25.linear_attn.in_proj_z",
+    "model.layers.25.linear_attn.out_proj",
+    "model.layers.25.mlp.shared_expert.down_proj",
+    "model.layers.25.mlp.shared_expert.gate_proj",
+    "model.layers.25.mlp.shared_expert.up_proj",
+    "model.layers.25.mlp.switch_mlp.down_proj",
+    "model.layers.25.mlp.switch_mlp.gate_proj",
+    "model.layers.25.mlp.switch_mlp.up_proj",
+    "model.layers.25.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.26.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.26.linear_attn.in_proj_a",
+    "model.layers.26.linear_attn.in_proj_b",
+    "model.layers.26.linear_attn.in_proj_qkv",
+    "model.layers.26.linear_attn.in_proj_z",
+    "model.layers.26.linear_attn.out_proj",
+    "model.layers.26.mlp.shared_expert.down_proj",
+    "model.layers.26.mlp.shared_expert.gate_proj",
+    "model.layers.26.mlp.shared_expert.up_proj",
+    "model.layers.26.mlp.switch_mlp.down_proj",
+    "model.layers.26.mlp.switch_mlp.gate_proj",
+    "model.layers.26.mlp.switch_mlp.up_proj",
+    "model.layers.26.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.27.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.27.mlp.shared_expert.down_proj",
+    "model.layers.27.mlp.shared_expert.gate_proj",
+    "model.layers.27.mlp.shared_expert.up_proj",
+    "model.layers.27.mlp.switch_mlp.down_proj",
+    "model.layers.27.mlp.switch_mlp.gate_proj",
+    "model.layers.27.mlp.switch_mlp.up_proj",
+    "model.layers.27.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.27.self_attn.indexer.index_qk_proj",
+    "model.layers.27.self_attn.k_proj",
+    "model.layers.27.self_attn.o_proj",
+    "model.layers.27.self_attn.q_proj",
+    "model.layers.27.self_attn.v_proj",
+    "model.layers.28.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.28.linear_attn.in_proj_a",
+    "model.layers.28.linear_attn.in_proj_b",
+    "model.layers.28.linear_attn.in_proj_qkv",
+    "model.layers.28.linear_attn.in_proj_z",
+    "model.layers.28.linear_attn.out_proj",
+    "model.layers.28.mlp.shared_expert.down_proj",
+    "model.layers.28.mlp.shared_expert.gate_proj",
+    "model.layers.28.mlp.shared_expert.up_proj",
+    "model.layers.28.mlp.switch_mlp.down_proj",
+    "model.layers.28.mlp.switch_mlp.gate_proj",
+    "model.layers.28.mlp.switch_mlp.up_proj",
+    "model.layers.28.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.29.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.29.linear_attn.in_proj_a",
+    "model.layers.29.linear_attn.in_proj_b",
+    "model.layers.29.linear_attn.in_proj_qkv",
+    "model.layers.29.linear_attn.in_proj_z",
+    "model.layers.29.linear_attn.out_proj",
+    "model.layers.29.mlp.shared_expert.down_proj",
+    "model.layers.29.mlp.shared_expert.gate_proj",
+    "model.layers.29.mlp.shared_expert.up_proj",
+    "model.layers.29.mlp.switch_mlp.down_proj",
+    "model.layers.29.mlp.switch_mlp.gate_proj",
+    "model.layers.29.mlp.switch_mlp.up_proj",
+    "model.layers.29.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.3.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.3.mlp.shared_expert.down_proj",
+    "model.layers.3.mlp.shared_expert.gate_proj",
+    "model.layers.3.mlp.shared_expert.up_proj",
+    "model.layers.3.mlp.switch_mlp.down_proj",
+    "model.layers.3.mlp.switch_mlp.gate_proj",
+    "model.layers.3.mlp.switch_mlp.up_proj",
+    "model.layers.3.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.3.self_attn.indexer.index_qk_proj",
+    "model.layers.3.self_attn.k_proj",
+    "model.layers.3.self_attn.o_proj",
+    "model.layers.3.self_attn.q_proj",
+    "model.layers.3.self_attn.v_proj",
+    "model.layers.30.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.30.linear_attn.in_proj_a",
+    "model.layers.30.linear_attn.in_proj_b",
+    "model.layers.30.linear_attn.in_proj_qkv",
+    "model.layers.30.linear_attn.in_proj_z",
+    "model.layers.30.linear_attn.out_proj",
+    "model.layers.30.mlp.shared_expert.down_proj",
+    "model.layers.30.mlp.shared_expert.gate_proj",
+    "model.layers.30.mlp.shared_expert.up_proj",
+    "model.layers.30.mlp.switch_mlp.down_proj",
+    "model.layers.30.mlp.switch_mlp.gate_proj",
+    "model.layers.30.mlp.switch_mlp.up_proj",
+    "model.layers.30.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.31.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.31.mlp.shared_expert.down_proj",
+    "model.layers.31.mlp.shared_expert.gate_proj",
+    "model.layers.31.mlp.shared_expert.up_proj",
+    "model.layers.31.mlp.switch_mlp.down_proj",
+    "model.layers.31.mlp.switch_mlp.gate_proj",
+    "model.layers.31.mlp.switch_mlp.up_proj",
+    "model.layers.31.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.31.self_attn.indexer.index_qk_proj",
+    "model.layers.31.self_attn.k_proj",
+    "model.layers.31.self_attn.o_proj",
+    "model.layers.31.self_attn.q_proj",
+    "model.layers.31.self_attn.v_proj",
+    "model.layers.32.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.32.linear_attn.in_proj_a",
+    "model.layers.32.linear_attn.in_proj_b",
+    "model.layers.32.linear_attn.in_proj_qkv",
+    "model.layers.32.linear_attn.in_proj_z",
+    "model.layers.32.linear_attn.out_proj",
+    "model.layers.32.mlp.shared_expert.down_proj",
+    "model.layers.32.mlp.shared_expert.gate_proj",
+    "model.layers.32.mlp.shared_expert.up_proj",
+    "model.layers.32.mlp.switch_mlp.down_proj",
+    "model.layers.32.mlp.switch_mlp.gate_proj",
+    "model.layers.32.mlp.switch_mlp.up_proj",
+    "model.layers.32.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.33.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.33.linear_attn.in_proj_a",
+    "model.layers.33.linear_attn.in_proj_b",
+    "model.layers.33.linear_attn.in_proj_qkv",
+    "model.layers.33.linear_attn.in_proj_z",
+    "model.layers.33.linear_attn.out_proj",
+    "model.layers.33.mlp.shared_expert.down_proj",
+    "model.layers.33.mlp.shared_expert.gate_proj",
+    "model.layers.33.mlp.shared_expert.up_proj",
+    "model.layers.33.mlp.switch_mlp.down_proj",
+    "model.layers.33.mlp.switch_mlp.gate_proj",
+    "model.layers.33.mlp.switch_mlp.up_proj",
+    "model.layers.33.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.34.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.34.linear_attn.in_proj_a",
+    "model.layers.34.linear_attn.in_proj_b",
+    "model.layers.34.linear_attn.in_proj_qkv",
+    "model.layers.34.linear_attn.in_proj_z",
+    "model.layers.34.linear_attn.out_proj",
+    "model.layers.34.mlp.shared_expert.down_proj",
+    "model.layers.34.mlp.shared_expert.gate_proj",
+    "model.layers.34.mlp.shared_expert.up_proj",
+    "model.layers.34.mlp.switch_mlp.down_proj",
+    "model.layers.34.mlp.switch_mlp.gate_proj",
+    "model.layers.34.mlp.switch_mlp.up_proj",
+    "model.layers.34.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.35.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.35.mlp.shared_expert.down_proj",
+    "model.layers.35.mlp.shared_expert.gate_proj",
+    "model.layers.35.mlp.shared_expert.up_proj",
+    "model.layers.35.mlp.switch_mlp.down_proj",
+    "model.layers.35.mlp.switch_mlp.gate_proj",
+    "model.layers.35.mlp.switch_mlp.up_proj",
+    "model.layers.35.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.35.self_attn.indexer.index_qk_proj",
+    "model.layers.35.self_attn.k_proj",
+    "model.layers.35.self_attn.o_proj",
+    "model.layers.35.self_attn.q_proj",
+    "model.layers.35.self_attn.v_proj",
+    "model.layers.36.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.36.linear_attn.in_proj_a",
+    "model.layers.36.linear_attn.in_proj_b",
+    "model.layers.36.linear_attn.in_proj_qkv",
+    "model.layers.36.linear_attn.in_proj_z",
+    "model.layers.36.linear_attn.out_proj",
+    "model.layers.36.mlp.shared_expert.down_proj",
+    "model.layers.36.mlp.shared_expert.gate_proj",
+    "model.layers.36.mlp.shared_expert.up_proj",
+    "model.layers.36.mlp.switch_mlp.down_proj",
+    "model.layers.36.mlp.switch_mlp.gate_proj",
+    "model.layers.36.mlp.switch_mlp.up_proj",
+    "model.layers.36.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.37.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.37.linear_attn.in_proj_a",
+    "model.layers.37.linear_attn.in_proj_b",
+    "model.layers.37.linear_attn.in_proj_qkv",
+    "model.layers.37.linear_attn.in_proj_z",
+    "model.layers.37.linear_attn.out_proj",
+    "model.layers.37.mlp.shared_expert.down_proj",
+    "model.layers.37.mlp.shared_expert.gate_proj",
+    "model.layers.37.mlp.shared_expert.up_proj",
+    "model.layers.37.mlp.switch_mlp.down_proj",
+    "model.layers.37.mlp.switch_mlp.gate_proj",
+    "model.layers.37.mlp.switch_mlp.up_proj",
+    "model.layers.37.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.38.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.38.linear_attn.in_proj_a",
+    "model.layers.38.linear_attn.in_proj_b",
+    "model.layers.38.linear_attn.in_proj_qkv",
+    "model.layers.38.linear_attn.in_proj_z",
+    "model.layers.38.linear_attn.out_proj",
+    "model.layers.38.mlp.shared_expert.down_proj",
+    "model.layers.38.mlp.shared_expert.gate_proj",
+    "model.layers.38.mlp.shared_expert.up_proj",
+    "model.layers.38.mlp.switch_mlp.down_proj",
+    "model.layers.38.mlp.switch_mlp.gate_proj",
+    "model.layers.38.mlp.switch_mlp.up_proj",
+    "model.layers.38.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.39.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.39.mlp.shared_expert.down_proj",
+    "model.layers.39.mlp.shared_expert.gate_proj",
+    "model.layers.39.mlp.shared_expert.up_proj",
+    "model.layers.39.mlp.switch_mlp.down_proj",
+    "model.layers.39.mlp.switch_mlp.gate_proj",
+    "model.layers.39.mlp.switch_mlp.up_proj",
+    "model.layers.39.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.39.self_attn.indexer.index_qk_proj",
+    "model.layers.39.self_attn.k_proj",
+    "model.layers.39.self_attn.o_proj",
+    "model.layers.39.self_attn.q_proj",
+    "model.layers.39.self_attn.v_proj",
+    "model.layers.4.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.4.linear_attn.in_proj_a",
+    "model.layers.4.linear_attn.in_proj_b",
+    "model.layers.4.linear_attn.in_proj_qkv",
+    "model.layers.4.linear_attn.in_proj_z",
+    "model.layers.4.linear_attn.out_proj",
+    "model.layers.4.mlp.shared_expert.down_proj",
+    "model.layers.4.mlp.shared_expert.gate_proj",
+    "model.layers.4.mlp.shared_expert.up_proj",
+    "model.layers.4.mlp.switch_mlp.down_proj",
+    "model.layers.4.mlp.switch_mlp.gate_proj",
+    "model.layers.4.mlp.switch_mlp.up_proj",
+    "model.layers.4.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.40.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.40.linear_attn.in_proj_a",
+    "model.layers.40.linear_attn.in_proj_b",
+    "model.layers.40.linear_attn.in_proj_qkv",
+    "model.layers.40.linear_attn.in_proj_z",
+    "model.layers.40.linear_attn.out_proj",
+    "model.layers.40.mlp.shared_expert.down_proj",
+    "model.layers.40.mlp.shared_expert.gate_proj",
+    "model.layers.40.mlp.shared_expert.up_proj",
+    "model.layers.40.mlp.switch_mlp.down_proj",
+    "model.layers.40.mlp.switch_mlp.gate_proj",
+    "model.layers.40.mlp.switch_mlp.up_proj",
+    "model.layers.40.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.41.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.41.linear_attn.in_proj_a",
+    "model.layers.41.linear_attn.in_proj_b",
+    "model.layers.41.linear_attn.in_proj_qkv",
+    "model.layers.41.linear_attn.in_proj_z",
+    "model.layers.41.linear_attn.out_proj",
+    "model.layers.41.mlp.shared_expert.down_proj",
+    "model.layers.41.mlp.shared_expert.gate_proj",
+    "model.layers.41.mlp.shared_expert.up_proj",
+    "model.layers.41.mlp.switch_mlp.down_proj",
+    "model.layers.41.mlp.switch_mlp.gate_proj",
+    "model.layers.41.mlp.switch_mlp.up_proj",
+    "model.layers.41.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.42.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.42.linear_attn.in_proj_a",
+    "model.layers.42.linear_attn.in_proj_b",
+    "model.layers.42.linear_attn.in_proj_qkv",
+    "model.layers.42.linear_attn.in_proj_z",
+    "model.layers.42.linear_attn.out_proj",
+    "model.layers.42.mlp.shared_expert.down_proj",
+    "model.layers.42.mlp.shared_expert.gate_proj",
+    "model.layers.42.mlp.shared_expert.up_proj",
+    "model.layers.42.mlp.switch_mlp.down_proj",
+    "model.layers.42.mlp.switch_mlp.gate_proj",
+    "model.layers.42.mlp.switch_mlp.up_proj",
+    "model.layers.42.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.43.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.43.mlp.shared_expert.down_proj",
+    "model.layers.43.mlp.shared_expert.gate_proj",
+    "model.layers.43.mlp.shared_expert.up_proj",
+    "model.layers.43.mlp.switch_mlp.down_proj",
+    "model.layers.43.mlp.switch_mlp.gate_proj",
+    "model.layers.43.mlp.switch_mlp.up_proj",
+    "model.layers.43.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.43.self_attn.indexer.index_qk_proj",
+    "model.layers.43.self_attn.k_proj",
+    "model.layers.43.self_attn.o_proj",
+    "model.layers.43.self_attn.q_proj",
+    "model.layers.43.self_attn.v_proj",
+    "model.layers.44.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.44.linear_attn.in_proj_a",
+    "model.layers.44.linear_attn.in_proj_b",
+    "model.layers.44.linear_attn.in_proj_qkv",
+    "model.layers.44.linear_attn.in_proj_z",
+    "model.layers.44.linear_attn.out_proj",
+    "model.layers.44.mlp.shared_expert.down_proj",
+    "model.layers.44.mlp.shared_expert.gate_proj",
+    "model.layers.44.mlp.shared_expert.up_proj",
+    "model.layers.44.mlp.switch_mlp.down_proj",
+    "model.layers.44.mlp.switch_mlp.gate_proj",
+    "model.layers.44.mlp.switch_mlp.up_proj",
+    "model.layers.44.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.45.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.45.linear_attn.in_proj_a",
+    "model.layers.45.linear_attn.in_proj_b",
+    "model.layers.45.linear_attn.in_proj_qkv",
+    "model.layers.45.linear_attn.in_proj_z",
+    "model.layers.45.linear_attn.out_proj",
+    "model.layers.45.mlp.shared_expert.down_proj",
+    "model.layers.45.mlp.shared_expert.gate_proj",
+    "model.layers.45.mlp.shared_expert.up_proj",
+    "model.layers.45.mlp.switch_mlp.down_proj",
+    "model.layers.45.mlp.switch_mlp.gate_proj",
+    "model.layers.45.mlp.switch_mlp.up_proj",
+    "model.layers.45.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.46.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.46.linear_attn.in_proj_a",
+    "model.layers.46.linear_attn.in_proj_b",
+    "model.layers.46.linear_attn.in_proj_qkv",
+    "model.layers.46.linear_attn.in_proj_z",
+    "model.layers.46.linear_attn.out_proj",
+    "model.layers.46.mlp.shared_expert.down_proj",
+    "model.layers.46.mlp.shared_expert.gate_proj",
+    "model.layers.46.mlp.shared_expert.up_proj",
+    "model.layers.46.mlp.switch_mlp.down_proj",
+    "model.layers.46.mlp.switch_mlp.gate_proj",
+    "model.layers.46.mlp.switch_mlp.up_proj",
+    "model.layers.46.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.47.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.47.mlp.shared_expert.down_proj",
+    "model.layers.47.mlp.shared_expert.gate_proj",
+    "model.layers.47.mlp.shared_expert.up_proj",
+    "model.layers.47.mlp.switch_mlp.down_proj",
+    "model.layers.47.mlp.switch_mlp.gate_proj",
+    "model.layers.47.mlp.switch_mlp.up_proj",
+    "model.layers.47.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.47.self_attn.indexer.index_qk_proj",
+    "model.layers.47.self_attn.k_proj",
+    "model.layers.47.self_attn.o_proj",
+    "model.layers.47.self_attn.q_proj",
+    "model.layers.47.self_attn.v_proj",
+    "model.layers.5.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.5.linear_attn.in_proj_a",
+    "model.layers.5.linear_attn.in_proj_b",
+    "model.layers.5.linear_attn.in_proj_qkv",
+    "model.layers.5.linear_attn.in_proj_z",
+    "model.layers.5.linear_attn.out_proj",
+    "model.layers.5.mlp.shared_expert.down_proj",
+    "model.layers.5.mlp.shared_expert.gate_proj",
+    "model.layers.5.mlp.shared_expert.up_proj",
+    "model.layers.5.mlp.switch_mlp.down_proj",
+    "model.layers.5.mlp.switch_mlp.gate_proj",
+    "model.layers.5.mlp.switch_mlp.up_proj",
+    "model.layers.5.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.6.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.6.linear_attn.in_proj_a",
+    "model.layers.6.linear_attn.in_proj_b",
+    "model.layers.6.linear_attn.in_proj_qkv",
+    "model.layers.6.linear_attn.in_proj_z",
+    "model.layers.6.linear_attn.out_proj",
+    "model.layers.6.mlp.shared_expert.down_proj",
+    "model.layers.6.mlp.shared_expert.gate_proj",
+    "model.layers.6.mlp.shared_expert.up_proj",
+    "model.layers.6.mlp.switch_mlp.down_proj",
+    "model.layers.6.mlp.switch_mlp.gate_proj",
+    "model.layers.6.mlp.switch_mlp.up_proj",
+    "model.layers.6.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.7.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.7.mlp.shared_expert.down_proj",
+    "model.layers.7.mlp.shared_expert.gate_proj",
+    "model.layers.7.mlp.shared_expert.up_proj",
+    "model.layers.7.mlp.switch_mlp.down_proj",
+    "model.layers.7.mlp.switch_mlp.gate_proj",
+    "model.layers.7.mlp.switch_mlp.up_proj",
+    "model.layers.7.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.7.self_attn.indexer.index_qk_proj",
+    "model.layers.7.self_attn.k_proj",
+    "model.layers.7.self_attn.o_proj",
+    "model.layers.7.self_attn.q_proj",
+    "model.layers.7.self_attn.v_proj",
+    "model.layers.8.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.8.linear_attn.in_proj_a",
+    "model.layers.8.linear_attn.in_proj_b",
+    "model.layers.8.linear_attn.in_proj_qkv",
+    "model.layers.8.linear_attn.in_proj_z",
+    "model.layers.8.linear_attn.out_proj",
+    "model.layers.8.mlp.shared_expert.down_proj",
+    "model.layers.8.mlp.shared_expert.gate_proj",
+    "model.layers.8.mlp.shared_expert.up_proj",
+    "model.layers.8.mlp.switch_mlp.down_proj",
+    "model.layers.8.mlp.switch_mlp.gate_proj",
+    "model.layers.8.mlp.switch_mlp.up_proj",
+    "model.layers.8.mlp_hyper_connection.input_mix_weight_down",
+    "model.layers.9.attn_hyper_connection.input_mix_weight_down",
+    "model.layers.9.linear_attn.in_proj_a",
+    "model.layers.9.linear_attn.in_proj_b",
+    "model.layers.9.linear_attn.in_proj_qkv",
+    "model.layers.9.linear_attn.in_proj_z",
+    "model.layers.9.linear_attn.out_proj",
+    "model.layers.9.mlp.shared_expert.down_proj",
+    "model.layers.9.mlp.shared_expert.gate_proj",
+    "model.layers.9.mlp.shared_expert.up_proj",
+    "model.layers.9.mlp.switch_mlp.down_proj",
+    "model.layers.9.mlp.switch_mlp.gate_proj",
+    "model.layers.9.mlp.switch_mlp.up_proj",
+    "model.layers.9.mlp_hyper_connection.input_mix_weight_down"
+  ],
+  "missing": [],
+  "mismatched": [],
+  "zero_count_experts": 63
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,305 @@
+{
+    "add_prefix_space": false,
+    "added_tokens_decoder": {
+        "248044": {
+            "content": "<|endoftext|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248045": {
+            "content": "<|im_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248046": {
+            "content": "<|im_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248047": {
+            "content": "<|object_ref_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248048": {
+            "content": "<|object_ref_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248049": {
+            "content": "<|box_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248050": {
+            "content": "<|box_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248051": {
+            "content": "<|quad_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248052": {
+            "content": "<|quad_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248053": {
+            "content": "<|vision_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248054": {
+            "content": "<|vision_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248055": {
+            "content": "<|vision_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248056": {
+            "content": "<|image_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248057": {
+            "content": "<|video_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248058": {
+            "content": "<tool_call>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248059": {
+            "content": "</tool_call>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248060": {
+            "content": "<|fim_prefix|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248061": {
+            "content": "<|fim_middle|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248062": {
+            "content": "<|fim_suffix|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248063": {
+            "content": "<|fim_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248064": {
+            "content": "<|repo_name|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248065": {
+            "content": "<|file_sep|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248066": {
+            "content": "<tool_response>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248067": {
+            "content": "</tool_response>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248068": {
+            "content": "<think>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248069": {
+            "content": "</think>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248070": {
+            "content": "<|audio_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248071": {
+            "content": "<|audio_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248072": {
+            "content": "<tts_pad>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248073": {
+            "content": "<tts_text_bos>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248074": {
+            "content": "<tts_text_eod>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248075": {
+            "content": "<tts_text_bos_single>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248076": {
+            "content": "<|audio_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        }
+    },
+    "additional_special_tokens": [
+        "<|im_start|>",
+        "<|im_end|>",
+        "<|object_ref_start|>",
+        "<|object_ref_end|>",
+        "<|box_start|>",
+        "<|box_end|>",
+        "<|quad_start|>",
+        "<|quad_end|>",
+        "<|vision_start|>",
+        "<|vision_end|>",
+        "<|vision_pad|>",
+        "<|image_pad|>",
+        "<|video_pad|>"
+    ],
+    "bos_token": null,
+    "chat_template": "{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- macro render_content(content, do_vision_count, is_system_content=false) %}\n    {%- if content is string %}\n        {{- content }}\n    {%- elif content is iterable and content is not mapping %}\n        {%- for item in content %}\n            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}\n                {%- if is_system_content %}\n                    {{- raise_exception('System message cannot contain images.') }}\n                {%- endif %}\n                {%- if do_vision_count %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                {%- endif %}\n                {%- if add_vision_id %}\n                    {{- 'Picture ' ~ image_count.value ~ ': ' }}\n                {%- endif %}\n                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}\n            {%- elif 'video' in item or item.type == 'video' %}\n                {%- if is_system_content %}\n                    {{- raise_exception('System message cannot contain videos.') }}\n                {%- endif %}\n                {%- if do_vision_count %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                {%- endif %}\n                {%- if add_vision_id %}\n                    {{- 'Video ' ~ video_count.value ~ ': ' }}\n                {%- endif %}\n                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}\n            {%- elif 'text' in item %}\n                {{- item.text }}\n            {%- else %}\n                {{- raise_exception('Unexpected item type in content.') }}\n            {%- endif %}\n        {%- endfor %}\n    {%- elif content is none or content is undefined %}\n        {{- '' }}\n    {%- else %}\n        {{- raise_exception('Unexpected content type.') }}\n    {%- endif %}\n{%- endmacro %}\n{%- if not messages %}\n    {{- raise_exception('No messages provided.') }}\n{%- endif %}\n{%- set reasoning_instructions = '' %}\n{%- if enable_thinking is undefined or enable_thinking is true %}\n    {%- set resolved_reasoning_effort = reasoning_effort|default('xhigh') %}\n    {%- if resolved_reasoning_effort not in ('xhigh', 'medium', 'low') %}\n        {{- raise_exception('Unexpected reasoning effort ' ~ reasoning_effort ~ '. Supported types are xhigh (default), medium, and low.') }}\n    {%- endif %}\n    {%- if resolved_reasoning_effort == 'xhigh' %}\n        {%- set reasoning_instructions = 'Reasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.' %}\n    {%- elif resolved_reasoning_effort == 'low' %}\n        {%- set reasoning_instructions = 'Reasoning effort is set to low. Keep your thinking brief and focused, moving directly to the conclusion without unnecessary elaboration.' %}\n    {%- endif %}\n{%- endif %}\n{%- if tools and tools is iterable and tools is not mapping %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if reasoning_instructions %}\n        {{- reasoning_instructions + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou have access to the following functions:\\n\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\" }}\n    {{- '\\n\\nIf you choose to call a function ONLY reply in the following format with NO suffix:\\n\\n<tool_call>\\n<function=example_function_name>\\n<parameter=example_parameter_1>\\nvalue_1\\n</parameter>\\n<parameter=example_parameter_2>\\nThis is the value for the second parameter\\nthat can span\\nmultiple lines\\n</parameter>\\n</function>\\n</tool_call>\\n\\n<IMPORTANT>\\nReminder:\\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\\n- Required parameters MUST be specified\\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\\n</IMPORTANT>' }}\n    {%- if messages[0].role == 'system' %}\n        {%- set content = render_content(messages[0].content, false, true)|trim %}\n        {%- if content %}\n            {{- '\\n\\n' + content }}\n        {%- endif %}\n    {%- endif %}\n    {{- '<|im_end|>\\n' }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {%- set content = render_content(messages[0].content, false, true)|trim %}\n        {%- if content %}\n            {{- '<|im_start|>system\\n' + (reasoning_instructions + '\\n\\n' if reasoning_instructions else '')  + content + '<|im_end|>\\n' }}\n        {%- elif reasoning_instructions %}\n            {{- '<|im_start|>system\\n' + reasoning_instructions + '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- elif reasoning_instructions %}\n        {{- '<|im_start|>system\\n' + reasoning_instructions + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" %}\n        {%- set content = render_content(message.content, false)|trim %}\n        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}\n            {%- set ns.multi_step_tool = false %}\n            {%- set ns.last_query_index = index %}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if ns.multi_step_tool %}\n    {{- raise_exception('No user query found in messages.') }}\n{%- endif %}\n{%- for message in messages %}\n    {%- set content = render_content(message.content, true)|trim %}\n    {%- if message.role == \"system\" %}\n        {%- if not loop.first %}\n            {{- raise_exception('System message must be at the beginning.') }}\n        {%- endif %}\n    {%- elif message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- endif %}\n        {%- set reasoning_content = reasoning_content|trim %}\n        {%- if preserve_thinking is undefined or preserve_thinking is true or loop.index0 > ns.last_query_index %}\n            {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content + '\\n</think>\\n\\n' + content }}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if tool_call.function is defined %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {%- if loop.first %}\n                    {%- if content|trim %}\n                        {{- '\\n\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n                    {%- else %}\n                        {{- '<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n                    {%- endif %}\n                {%- else %}\n                    {{- '\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n                {%- endif %}\n                {%- if tool_call.arguments is defined and tool_call.arguments != '' %}\n                    {%- for args_name, args_value in tool_call.arguments|items %}\n                        {{- '<parameter=' + args_name + '>\\n' }}\n                        {%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}\n                        {{- args_value }}\n                        {{- '\\n</parameter>\\n' }}\n                    {%- endfor %}\n                {%- endif %}\n                {{- '</function>\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.previtem and loop.previtem.role != \"tool\" %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if not loop.last and loop.nextitem.role != \"tool\" %}\n            {{- '<|im_end|>\\n' }}\n        {%- elif loop.last %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- else %}\n        {{- raise_exception('Unexpected message role.') }}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- else %}\n        {{- '<think>\\n' }}\n    {%- endif %}\n{%- endif %}",
+    "clean_up_tokenization_spaces": false,
+    "eos_token": "<|im_end|>",
+    "errors": "replace",
+    "model_max_length": 262144,
+    "pad_token": "<|endoftext|>",
+    "split_special_tokens": false,
+    "tokenizer_class": "Qwen2Tokenizer",
+    "unk_token": null,
+    "add_bos_token": false,
+    "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+    "extra_special_tokens": {
+        "audio_bos_token": "<|audio_start|>",
+        "audio_eos_token": "<|audio_end|>",
+        "audio_token": "<|audio_pad|>",
+        "image_token": "<|image_pad|>",
+        "video_token": "<|video_pad|>",
+        "vision_bos_token": "<|vision_start|>",
+        "vision_eos_token": "<|vision_end|>"
+    }
+}