Duplicate from Vontra/MiMo-V2.6-Flash-RL-MLX-4bit-MTP

Co-authored-by: Ashley Hart <ashxhart@users.noreply.huggingface.co>

.gitattributes

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

README.md

@@ -0,0 +1,105 @@
+---
+license: mit
+language:
+- en
+- zh
+library_name: mlx
+pipeline_tag: text-generation
+base_model: XiaomiMiMo/MiMo-V2.6-Flash-RL
+base_model_relation: quantized
+tags:
+- mlx
+- apple-silicon
+- mimo-v2
+- mixture-of-experts
+- 4-bit
+- mtp
+---
+
+<div align="center">
+  <picture>
+    <source srcset="https://github.com/XiaomiMiMo/MiMo/raw/main/figures/Xiaomi_MiMo_darkmode.png?raw=true" media="(prefers-color-scheme: dark)">
+    <img src="https://github.com/XiaomiMiMo/MiMo/raw/main/figures/Xiaomi_MiMo.png?raw=true" width="48%" alt="Xiaomi MiMo">
+  </picture>
+</div>
+
+<h1 align="center">MiMo-V2.6-Flash-RL MLX 4-bit MTP</h1>
+
+<p align="center">
+  A tested Apple Silicon conversion of
+  <a href="https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL">XiaomiMiMo/MiMo-V2.6-Flash-RL</a>,
+  published by <a href="https://huggingface.co/Vontra">Vontra</a>.
+</p>
+
+## What is in this release
+
+This is the text backbone of MiMo-V2.6-Flash-RL converted for MLX. The dense projections use 4-bit affine quantization with group size 64, while the model's native MXFP4 MoE experts remain in their original group-size-32 format. The resulting main model averages 4.257 bits per weight.
+
+The checkpoint includes its native three-layer MTP payload in `mtp/model_mtp.safetensors`, converted to 4-bit affine weights. It also carries the upstream five-layer DFlash drafter, vision encoder, audio encoder, and audio tokenizer so those assets do not need a second download.
+
+| Native component | Path | Format |
+| --- | --- | --- |
+| MTP predictor | `mtp/model_mtp.safetensors` | MLX 4-bit affine |
+| DFlash drafter | `dflash/model.safetensors` | Upstream BF16 |
+| Vision encoder | `omnimodal/vision_encoder.safetensors` | Upstream BF16 |
+| Audio encoder | `omnimodal/audio_encoder.safetensors` | Upstream BF16 |
+| Audio tokenizer | `audio_tokenizer/model.safetensors` | Upstream weights |
+
+Current oMLX and MLX text generation run the target model correctly but do not automatically execute MiMo's MTP, DFlash, vision, or audio paths. The speed figures below are serial text decode measurements. The auxiliary tensors and configs are packaged for MiMo-aware runtimes and ongoing MLX integration, not advertised as working oMLX controls.
+
+## Measured on Apple Silicon
+
+Tested on a 256 GB M3 Ultra Mac Studio with oMLX 0.7.0.dev2, MLX 0.32.2, and mlx-lm 0.31.3.
+
+| Test | Result |
+| --- | ---: |
+| Sustained generation, 128-token decode | 59.4 tok/s average |
+| Prompt processing, 512 tokens | 477.1 tok/s |
+| Prompt processing, 2,048 tokens | 562.8 tok/s |
+| Peak unified memory, short context | 164.3 GB |
+| Peak unified memory, 2,048-token prompt | 166.8 GB |
+| Quantized text model size on disk | about 154 GiB |
+| Complete repository size | about 160 GiB |
+
+The model produced correct arithmetic, a clear factual explanation, and coherent Python in repeated smoke tests. A 256 GB Mac is recommended so there is room for the model, KV cache, and the rest of the system.
+
+## Run with mlx-lm
+
+```bash
+pip install -U "mlx-lm>=0.31.3"
+
+python -m mlx_lm generate \
+  --model Vontra/MiMo-V2.6-Flash-RL-MLX-4bit-MTP \
+  --prompt "Write a Python function that checks whether an integer is prime." \
+  --max-tokens 256 \
+  --temp 0.6
+```
+
+The upstream tokenizer chat template is included. In oMLX, download or select `Vontra/MiMo-V2.6-Flash-RL-MLX-4bit-MTP` as an MLX model.
+
+## Quantization notes
+
+MiMo-V2.6 stores fused attention tensors in checkpoint tensor-parallel order and pads the FP8 scale grid separately for each shard. This conversion reconstructs those shards before quantization. Skipping that step produces a model that loads but returns broken output.
+
+The quantized MTP payload lives in its own `mtp/` directory with a manifest describing its tensors. It is not a standalone drafter for `mlx_lm.generate --draft-model` today. The upstream DFlash payload retains its trained mask embedding and corrected JSON config; an MLX smoke test matched serial greedy output, but it did not beat serial decode in the current experimental runtime.
+
+The vision and audio tensors are kept outside the root text-model index so `mlx_lm` and oMLX continue to load the tested text checkpoint unchanged. `omnimodal/manifest.json` records every auxiliary path and its upstream source.
+
+## About MiMo-V2.6-Flash-RL
+
+Xiaomi describes MiMo-V2.6-Flash-RL as a sparse 309B-parameter MoE with 15B active parameters, 48 transformer layers, 256 routed experts, and eight active experts per token. The full upstream release supports a one-million-token context and omnimodal inputs. See the [original model card](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL) for architecture details, evaluations, deployment recipes, intended use, and limitations.
+
+## License and credit
+
+The upstream model is released under the MIT license. All model architecture, training, tokenizer work, and original branding belong to the Xiaomi MiMo team. This repository contains a community MLX conversion and measured Apple Silicon results.
+
+```bibtex
+@misc{mimo2026v26flash,
+  title={MiMo-V2.6-Flash-RL},
+  author={{Xiaomi MiMo Team}},
+  year={2026},
+  howpublished={\url{https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL}},
+}
+```
+
+Follow [Vontra](https://huggingface.co/Vontra) for new Apple Silicon releases and fixes.

audio_tokenizer/chat_template.jinja

@@ -0,0 +1,120 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {%- if messages[0].content is string %}
+            {{- messages[0].content }}
+        {%- else %}
+            {%- for content in messages[0].content %}
+                {%- if content.type == 'audio' %}
+                    {{- ("<|sosp|>" + (content.meta | tojson) + "<|eosp|>") }}
+                {%- elif content.type == 'text' %}
+                    {{- content.text }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+    {%- endif %}
+    {{- '\n\n' }}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' }}
+        {%- if messages[0].content is string %}
+            {{- messages[0].content }}
+        {%- else %}
+            {%- for content in messages[0].content %}
+                {%- if content.type == 'audio' %}
+                    {{- ("<|sosp|>" + (content.meta | tojson) + "<|eosp|>") }}
+                {%- elif content.type == 'text' %}
+                    {{- content.text }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '\n<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1, assistant_is_last=false) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = namespace(text="") %}
+        {%- for mcontent in message.content %}
+            {%- if mcontent.type == 'audio' %}
+                {%- set content.text = content.text~("<|sosp|>" + (mcontent.meta | tojson) + "<|eosp|>") %}
+            {%- elif mcontent.type == 'text' %}
+                {%- set content.text = content.text~mcontent.text %}
+            {%- endif %}
+        {%- endfor %}
+        {%- set content = content.text %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|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 %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip("\n") + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}{{- '\n' }}{%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {%- if loop.last %}
+            {%- set ns.assistant_is_last = true %}
+        {%- else %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}{{- '<|im_start|>user' }}{%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}{{- '<|im_end|>\n' }}{%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt and not ns.assistant_is_last %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if audio_output  %}
+        {{- '<|sostm|>'}}
+    {%- elif not enable_thinking %}
+        {{- '<think>\n\n</think>\n' }}
+    {%- endif %}
+{%- endif %}

audio_tokenizer/config.json

@@ -0,0 +1,68 @@
+{
+  "max_audio_seconds": 300,
+  "stride_size": 2,
+  "avg_pooler": 2,
+  "d_model": 1024,
+  "scale_embedding": false,
+  "kernel_size": 3,
+  "activation_function": "gelu",
+  "encoder_layers": 24,
+  "encoder_skip_layer_id": 3,
+  "encoder_attention_heads": 16,
+  "encoder_ffn_dim": 4096,
+  "encoder_causal": true,
+  "encoder_attn_window_size": [
+    128,
+    0
+  ],
+  "decoder_layers": 24,
+  "decoder_attention_heads": 16,
+  "decoder_ffn_dim": 4096,
+  "decoder_kernel_size": 3,
+  "decoder_stride_size": 2,
+  "decoder_causal": true,
+  "decoder_attn_window_size": [
+    128,
+    0
+  ],
+  "nfft": 960,
+  "n_mels": 128,
+  "sampling_rate": 24000,
+  "hop_length": 240,
+  "window_size": 960,
+  "vocoder_padding": "same",
+  "fmin": 0,
+  "fmax": null,
+  "num_quantizers": 20,
+  "codebook_size": [
+    1024,
+    1024,
+    256,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128,
+    128
+  ],
+  "threshold_ema_dead_code": 2,
+  "position_embedding_type": "rope",
+  "rope_theta": 10000,
+  "rope_type": "default",
+  "ln_type": "LayerNorm",
+  "use_istft_only": true,
+  "hybrid_attention": true,
+  "hybrid_block_size": 8,
+  "swa_per_block": 2
+}

audio_tokenizer/generation_config.json

@@ -0,0 +1,9 @@
+{
+  "do_sample": true,
+  "temperature": 0.6,
+  "top_k": -1,
+  "top_p": 0.95,
+  "audio_temperature": 0.9,
+  "audio_top_k": -1,
+  "audio_top_p": 0.95
+}

audio_tokenizer/model.safetensors

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

audio_tokenizer/tokenizer_config.json

@@ -0,0 +1,267 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<|mimo_audio_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151666": {
+      "content": "<|mimo_audio_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151669": {
+      "content": "<|audio_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151670": {
+      "content": "<|mimo_video_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151671": {
+      "content": "<|mimo_video_end|>",
+      "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|>",
+    "<|audio_pad|>",
+    "<|mimo_audio_start|>",
+    "<|mimo_audio_end|>",
+    "<|mimo_video_start|>",
+    "<|mimo_video_end|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

chat_template.jinja

@@ -0,0 +1,96 @@
+{%- macro render_value(value) -%}
+    {%- if value is string -%}
+        {{- value -}}
+    {%- else -%}
+        {{- value | tojson(ensure_ascii=False) -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_content(message_content) -%}
+    {%- if message_content is string -%}
+        {{- message_content -}}
+    {%- elif message_content is iterable -%}
+        {%- for part in message_content -%}
+            {%- if part is not mapping -%}
+                {{- part -}}
+            {%- elif part['type'] == 'image' or 'image' in part or 'image_url' in part -%}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' -}}
+            {%- elif part['type'] == 'audio' or part['type'] == 'input_audio' or 'audio' in part or 'audio_url' in part or 'input_audio' in part -%}
+                {{- '<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>' -}}
+            {%- elif part['type'] == 'video' or 'video' in part or 'video_url' in part -%}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' -}}
+            {%- elif 'text' in part -%}
+                {{- part['text'] -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tools(tools) -%}
+    {{- 'You are provided with the following tools:\n\n<tools>' -}}
+    {%- for tool in tools -%}
+        {{- '\n' ~ (tool | tojson(ensure_ascii=False)) -}}
+    {%- endfor -%}
+    {{- '\n</tools>' -}}
+{%- endmacro -%}
+
+{%- macro render_tool_calls(tool_calls) -%}
+    {%- for tool_call in tool_calls -%}
+        {%- if tool_call.function is defined -%}
+            {%- set tool_call = tool_call.function -%}
+        {%- elif tool_call.custom is defined -%}
+            {%- set tool_call = tool_call.custom -%}
+        {%- endif -%}
+        {{- '<tool_call><function=' ~ tool_call.name ~ '>' -}}
+        {%- if tool_call.input is defined and tool_call.input is string -%}
+            {{- tool_call.input -}}
+        {%- elif tool_call.arguments -%}
+            {%- if tool_call.arguments is string -%}
+                {{- tool_call.arguments -}}
+            {%- else -%}
+                {%- for args_name, args_value in tool_call.arguments | items -%}
+                    {{- '<parameter=' ~ args_name ~ '>' ~ render_value(args_value) ~ '</parameter>' -}}
+                {%- endfor -%}
+            {%- endif -%}
+        {%- endif -%}
+        {{- '</function></tool_call>' -}}
+    {%- endfor -%}
+{%- endmacro -%}
+
+{%- macro render_assistant_message(message) -%}
+    {%- set content = render_content(message.content) -%}
+    {%- set reasoning = message.reasoning_content if message.reasoning_content is string else '' -%}
+    {{- '<|im_start|>assistant\n<think>' ~ reasoning ~ '</think>' ~ content -}}
+    {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 -%}
+        {{- render_tool_calls(message.tool_calls) -}}
+    {%- endif -%}
+    {{- '<|im_end|>' -}}
+{%- endmacro -%}
+
+
+{%- if tools is defined and tools is iterable and tools | length > 0 -%}
+    {{- '<|im_start|>system\n' ~ render_tools(tools) ~ '<|im_end|>' -}}
+{%- endif -%}
+
+{%- for message in messages -%}
+    {%- if message.role == 'assistant' -%}
+        {{- render_assistant_message(message) -}}
+    {%- else -%}
+        {%- set body = render_content(message.content) -%}
+        {{- '<|im_start|>' ~ message.role ~ '\n' ~ body -}}
+        {%- if message.tools is defined and message.tools is iterable and message.tools | length > 0 -%}
+            {%- if body -%}
+                {{- '\n\n' -}}
+            {%- endif -%}
+            {{- render_tools(message.tools) -}}
+        {%- endif -%}
+        {{- '<|im_end|>' -}}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {{- '<|im_start|>assistant\n' -}}
+    {%- if enable_thinking is false -%}
+        {{- '<think></think>' -}}
+    {%- endif -%}
+{%- endif -%}

configuration_mimo_v2.py

@@ -0,0 +1,247 @@
+# coding=utf-8
+#
+# Copyright 2026 Xiaomi Corporation.
+# Copyright 2026 The HuggingFace Inc. team.
+#
+# 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.
+
+from copy import deepcopy
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+_MIMOV2_ATTENTION_PROJECTION_LAYOUTS = {"split", "fused_qkv"}
+
+_MIMOV2_SPLIT_TP_PLAN = {
+    "layers.*.self_attn.q_proj": "colwise",
+    "layers.*.self_attn.k_proj": "colwise",
+    "layers.*.self_attn.v_proj": "colwise",
+    "layers.*.self_attn.o_proj": "rowwise",
+    "layers.*.mlp.gate_proj": "colwise",
+    "layers.*.mlp.up_proj": "colwise",
+    "layers.*.mlp.down_proj": "rowwise",
+}
+
+_MIMOV2_FUSED_QKV_TP_PLAN = {
+    "layers.*.self_attn.qkv_proj": "colwise",
+    "layers.*.self_attn.o_proj": "rowwise",
+    "layers.*.mlp.gate_proj": "colwise",
+    "layers.*.mlp.up_proj": "colwise",
+    "layers.*.mlp.down_proj": "rowwise",
+}
+
+_MIMOV2_PP_PLAN = {
+    "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+    "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+    "norm": (["hidden_states"], ["hidden_states"]),
+}
+
+
+def _to_plain_dict(value):
+    if value is None:
+        return {}
+    if isinstance(value, dict):
+        return deepcopy(value)
+    if hasattr(value, "to_dict"):
+        return deepcopy(value.to_dict())
+    if hasattr(value, "__dict__"):
+        return deepcopy(vars(value))
+    raise TypeError(f"Unsupported config value type: {type(value)!r}")
+
+
+class MiMoV2Config(PretrainedConfig):
+
+    model_type = "mimo_v2"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    base_model_tp_plan = _MIMOV2_SPLIT_TP_PLAN
+    base_model_pp_plan = _MIMOV2_PP_PLAN
+
+    attribute_map = {
+        "num_local_experts": "n_routed_experts",
+    }
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=4096,
+        intermediate_size=22016,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=32,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        layernorm_epsilon=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_dropout=0.0,
+        attention_bias=False,
+        attention_value_scale=None,
+        head_dim=None,
+        v_head_dim=None,
+        swa_num_attention_heads=None,
+        swa_num_key_value_heads=None,
+        swa_head_dim=None,
+        swa_v_head_dim=None,
+        swa_rope_theta=None,
+        sliding_window=None,
+        sliding_window_size=None,
+        add_full_attention_sink_bias=False,
+        add_swa_attention_sink_bias=False,
+        hybrid_block_size=None,
+        hybrid_layer_pattern=None,
+        partial_rotary_factor=1.0,
+        n_routed_experts=None,
+        moe_intermediate_size=None,
+        num_experts_per_tok=None,
+        routed_scaling_factor=None,
+        scoring_func="sigmoid",
+        topk_method="noaux_tc",
+        n_group=None,
+        topk_group=None,
+        norm_topk_prob=True,
+        moe_layer_freq=None,
+        attention_projection_layout="split",
+        vision_config=None,
+        audio_config=None,
+        processor_config=None,
+        image_token_id=None,
+        video_token_id=None,
+        vision_start_token_id=None,
+        vision_end_token_id=None,
+        vision_model_type=None,
+        **kwargs,
+    ):
+        rope_parameters = kwargs.pop("rope_parameters", None)
+        if rope_scaling is None and rope_parameters is not None:
+            rope_scaling = rope_parameters
+
+        if attention_projection_layout is None:
+            attention_projection_layout = "split"
+        if attention_projection_layout not in _MIMOV2_ATTENTION_PROJECTION_LAYOUTS:
+            raise ValueError(f"Unsupported MiMoV2 attention projection layout: {attention_projection_layout}")
+
+        self.attention_projection_layout = attention_projection_layout
+        self.base_model_tp_plan = (
+            _MIMOV2_FUSED_QKV_TP_PLAN.copy()
+            if attention_projection_layout == "fused_qkv"
+            else _MIMOV2_SPLIT_TP_PLAN.copy()
+        )
+        self.base_model_pp_plan = _MIMOV2_PP_PLAN.copy()
+
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        if num_attention_heads % num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.layernorm_epsilon = layernorm_epsilon
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_dropout = attention_dropout
+        self.attention_bias = attention_bias
+        self.attention_value_scale = attention_value_scale
+
+        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
+        self.v_head_dim = v_head_dim if v_head_dim is not None else self.head_dim
+        self.swa_num_attention_heads = (
+            swa_num_attention_heads if swa_num_attention_heads is not None else num_attention_heads
+        )
+        self.swa_num_key_value_heads = (
+            swa_num_key_value_heads if swa_num_key_value_heads is not None else num_key_value_heads
+        )
+        if self.swa_num_attention_heads % self.swa_num_key_value_heads != 0:
+            raise ValueError("swa_num_attention_heads must be divisible by swa_num_key_value_heads")
+        self.swa_head_dim = swa_head_dim if swa_head_dim is not None else self.head_dim
+        self.swa_v_head_dim = swa_v_head_dim if swa_v_head_dim is not None else self.swa_head_dim
+        self.swa_rope_theta = swa_rope_theta if swa_rope_theta is not None else rope_theta
+
+        if sliding_window is None:
+            sliding_window = sliding_window_size
+        self.sliding_window = sliding_window
+        self.sliding_window_size = sliding_window_size if sliding_window_size is not None else sliding_window
+        self.add_full_attention_sink_bias = add_full_attention_sink_bias
+        self.add_swa_attention_sink_bias = add_swa_attention_sink_bias
+
+        if hybrid_block_size is not None and hybrid_layer_pattern is None:
+            hybrid_layer_pattern = [0 if ((i + 1) % hybrid_block_size == 0) else 1 for i in range(num_hidden_layers)]
+        elif hybrid_layer_pattern is None:
+            hybrid_layer_pattern = [0] * num_hidden_layers
+        if len(hybrid_layer_pattern) != num_hidden_layers:
+            raise ValueError("hybrid_layer_pattern length must match num_hidden_layers")
+        self.hybrid_block_size = hybrid_block_size
+        self.hybrid_layer_pattern = hybrid_layer_pattern
+
+        self.partial_rotary_factor = partial_rotary_factor
+
+        self.n_routed_experts = n_routed_experts
+        self.moe_intermediate_size = moe_intermediate_size if moe_intermediate_size is not None else intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.routed_scaling_factor = routed_scaling_factor
+        self.scoring_func = scoring_func
+        self.topk_method = topk_method
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.norm_topk_prob = norm_topk_prob
+        if isinstance(moe_layer_freq, int):
+            moe_layer_freq = [moe_layer_freq > 0 and i % moe_layer_freq == 0 for i in range(num_hidden_layers)]
+        elif moe_layer_freq is None:
+            moe_layer_freq = [False] * num_hidden_layers
+        if len(moe_layer_freq) != num_hidden_layers:
+            raise ValueError("moe_layer_freq length must match num_hidden_layers")
+        self.moe_layer_freq = moe_layer_freq
+
+        self.vision_config = _to_plain_dict(vision_config)
+        self.audio_config = _to_plain_dict(audio_config)
+        self.processor_config = _to_plain_dict(processor_config)
+        self.image_token_id = image_token_id
+        self.video_token_id = video_token_id
+        self.vision_start_token_id = vision_start_token_id
+        self.vision_end_token_id = vision_end_token_id
+        self.vision_model_type = vision_model_type
+        self.audio_token_id = self.processor_config.get("audio_token_id", None) if self.processor_config else None
+        self.audio_start_token_id = (
+            self.processor_config.get("audio_start_token_id", None) if self.processor_config else None
+        )
+        self.audio_end_token_id = (
+            self.processor_config.get("audio_end_token_id", None) if self.processor_config else None
+        )
+
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+__all__ = ["MiMoV2Config"]

dflash/config.json

@@ -0,0 +1,56 @@
+{
+  "architectures": [
+    "DFlashDraftModel"
+  ],
+  "model_type": "qwen3",
+  "auto_map": {
+    "AutoModel": "dflash.DFlashDraftModel"
+  },
+  "hidden_size": 4096,
+  "intermediate_size": 16384,
+  "num_hidden_layers": 5,
+  "num_attention_heads": 64,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "v_head_dim": 128,
+  "partial_rotary_factor": 0.5,
+  "block_size": 8,
+  "dflash_config": {
+    "target_layer_ids": [
+      0,
+      11,
+      23,
+      35,
+      47
+    ],
+    "mask_token_id": 151675,
+    "num_anchors": 4096,
+    "block_size": 8,
+    "loss_decay_gamma": 7.0,
+    "attention_value_scale": 0.612,
+    "attention_sink_bias": true
+  },
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention"
+  ],
+  "sliding_window": 1024,
+  "use_sliding_window": true,
+  "is_causal": false,
+  "num_target_layers": 48,
+  "target_hidden_size": 4096,
+  "vocab_size": 152576,
+  "max_position_embeddings": 1048576,
+  "rope_theta": 10000.0,
+  "rms_norm_eps": 1e-06,
+  "torch_dtype": "bfloat16",
+  "hidden_act": "silu",
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "add_swa_attention_sink_bias": true,
+  "tie_word_embeddings": false,
+  "use_cache": true
+}

dflash/dflash.py

@@ -0,0 +1,379 @@
+from typing import Callable, Optional
+
+import torch
+from torch import nn
+from transformers import DynamicCache
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.qwen3.modeling_qwen3 import (
+    ALL_ATTENTION_FUNCTIONS,
+    FlashAttentionKwargs,
+    GradientCheckpointingLayer,
+    Qwen3Config,
+    Qwen3MLP,
+    Qwen3PreTrainedModel,
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    eager_attention_forward,
+    rotate_half,
+)
+from typing_extensions import Tuple, Unpack
+
+
+def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor:
+    if temperature < 1e-5:
+        return torch.argmax(logits, dim=-1)
+    bsz, seq_len, vocab_size = logits.shape
+    logits = logits.view(-1, vocab_size)
+    logits = logits / temperature
+    probs = torch.softmax(logits, dim=-1)
+    return torch.multinomial(probs, num_samples=1).view(bsz, seq_len)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_len = q.size(-2)
+    q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :])
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class Qwen3DFlashAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = False
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+        )
+        self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.sliding_window = (
+            config.sliding_window
+            if config.layer_types[layer_idx] == "sliding_attention"
+            else None
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        target_hidden: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        bsz, q_len = hidden_states.shape[:-1]
+        ctx_len = target_hidden.shape[1]
+        q = self.q_proj(hidden_states)
+        q = q.view(bsz, q_len, -1, self.head_dim)
+        q = self.q_norm(q).transpose(1, 2)
+        k_ctx = self.k_proj(target_hidden)
+        k_noise = self.k_proj(hidden_states)
+        v_ctx = self.v_proj(target_hidden)
+        v_noise = self.v_proj(hidden_states)
+        k = torch.cat([k_ctx, k_noise], dim=1).view(
+            bsz, ctx_len + q_len, -1, self.head_dim
+        )
+        v = torch.cat([v_ctx, v_noise], dim=1).view(
+            bsz, ctx_len + q_len, -1, self.head_dim
+        )
+        k = self.k_norm(k).transpose(1, 2)
+        v = v.transpose(1, 2)
+        cos, sin = position_embeddings
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs)
+        attn_fn: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        attn_output, attn_weights = attn_fn(
+            self,
+            q,
+            k,
+            v,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx)
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        target_hidden: Optional[torch.Tensor] = None,
+        hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[
+            Tuple[torch.Tensor, torch.Tensor]
+        ] = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            target_hidden=target_hidden,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )[0]
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+def build_target_layer_ids(num_target_layers: int, num_draft_layers: int):
+    if num_draft_layers == 1:
+        return [(num_target_layers // 2)]
+    start = 1
+    end = num_target_layers - 3
+    span = end - start
+    target_layer_ids = [
+        int(round(start + (i * span) / (num_draft_layers - 1)))
+        for i in range(num_draft_layers)
+    ]
+    return target_layer_ids
+
+
+def extract_context_feature(
+    hidden_states: list[torch.Tensor],
+    layer_ids: Optional[list[int]],
+) -> torch.Tensor:
+    offset = 1
+    selected_states = []
+    for layer_id in layer_ids:
+        selected_states.append(hidden_states[layer_id + offset])
+    target_hidden = torch.cat(selected_states, dim=-1)
+    return target_hidden
+
+
+class DFlashDraftModel(Qwen3PreTrainedModel):
+    config_class = Qwen3Config
+    _no_split_modules = ["Qwen3DFlashDecoderLayer"]
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.config = config
+        self.layers = nn.ModuleList(
+            [
+                Qwen3DFlashDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        dflash_config = getattr(config, "dflash_config", {}) or {}
+        self.target_layer_ids = dflash_config.get(
+            "target_layer_ids",
+            build_target_layer_ids(config.num_target_layers, config.num_hidden_layers),
+        )
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3RotaryEmbedding(config)
+        self.fc = nn.Linear(
+            len(self.target_layer_ids) * config.hidden_size,
+            config.hidden_size,
+            bias=False,
+        )
+        self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.block_size = config.block_size
+        self.mask_token_id = dflash_config.get("mask_token_id", None)
+        self.post_init()
+
+    def forward(
+        self,
+        position_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        noise_embedding: Optional[torch.Tensor] = None,
+        target_hidden: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_states = noise_embedding
+        target_hidden = self.hidden_norm(self.fc(target_hidden))
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states=hidden_states,
+                target_hidden=target_hidden,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                use_cache=use_cache,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+        return self.norm(hidden_states)
+
+    @torch.inference_mode()
+    def spec_generate(
+        self,
+        target: nn.Module,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int,
+        stop_token_ids: list[int],
+        temperature: float,
+    ):
+        self.eval()
+        num_input_tokens = input_ids.shape[1]
+        max_length = num_input_tokens + max_new_tokens
+
+        block_size = self.block_size
+        output_ids = torch.full(
+            (1, max_length + block_size),
+            self.mask_token_id,
+            dtype=torch.long,
+            device=target.device,
+        )
+        position_ids = torch.arange(
+            output_ids.shape[1], device=target.device
+        ).unsqueeze(0)
+
+        past_key_values_target = DynamicCache()
+        past_key_values_draft = DynamicCache()
+
+        # Prefill stage
+        output = target(
+            input_ids,
+            position_ids=position_ids[:, :num_input_tokens],
+            past_key_values=past_key_values_target,
+            use_cache=True,
+            logits_to_keep=1,
+            output_hidden_states=True,
+        )
+
+        output_ids[:, :num_input_tokens] = input_ids
+        output_ids[:, num_input_tokens : num_input_tokens + 1] = sample(
+            output.logits, temperature
+        )
+        target_hidden = extract_context_feature(
+            output.hidden_states, self.target_layer_ids
+        )
+
+        # Decode stage
+        acceptance_lengths = []
+        start = input_ids.shape[1]
+        while start < max_length:
+            block_output_ids = output_ids[:, start : start + block_size].clone()
+            block_position_ids = position_ids[:, start : start + block_size]
+            noise_embedding = target.model.embed_tokens(block_output_ids)
+            draft_logits = target.lm_head(
+                self(
+                    target_hidden=target_hidden,
+                    noise_embedding=noise_embedding,
+                    position_ids=position_ids[
+                        :, past_key_values_draft.get_seq_length() : start + block_size
+                    ],
+                    past_key_values=past_key_values_draft,
+                    use_cache=True,
+                    is_causal=False,
+                )[:, -block_size + 1 :, :]
+            )
+            past_key_values_draft.crop(start)
+            block_output_ids[:, 1:] = sample(draft_logits)
+
+            output = target(
+                block_output_ids,
+                position_ids=block_position_ids,
+                past_key_values=past_key_values_target,
+                use_cache=True,
+                output_hidden_states=True,
+            )
+
+            posterior = sample(output.logits, temperature)
+            acceptance_length = (
+                (block_output_ids[:, 1:] == posterior[:, :-1])
+                .cumprod(dim=1)
+                .sum(dim=1)[0]
+                .item()
+            )
+            output_ids[:, start : start + acceptance_length + 1] = block_output_ids[
+                :, : acceptance_length + 1
+            ]
+            output_ids[:, start + acceptance_length + 1] = posterior[
+                :, acceptance_length
+            ]
+            start += acceptance_length + 1
+            past_key_values_target.crop(start)
+            target_hidden = extract_context_feature(
+                output.hidden_states, self.target_layer_ids
+            )[:, : acceptance_length + 1, :]
+            acceptance_lengths.append(acceptance_length + 1)
+            if stop_token_ids is not None and any(
+                stop_token_id in output_ids[:, num_input_tokens:]
+                for stop_token_id in stop_token_ids
+            ):
+                break
+        output_ids = output_ids[:, :max_length]
+        output_ids = output_ids[:, output_ids[0] != self.mask_token_id]
+        if stop_token_ids is not None:
+            stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device)
+            stop_token_indices = torch.isin(
+                output_ids[0][num_input_tokens:], stop_token_ids
+            ).nonzero(as_tuple=True)[0]
+            if stop_token_indices.numel() > 0:
+                output_ids = output_ids[
+                    :, : num_input_tokens + stop_token_indices[0] + 1
+                ]
+
+        return output_ids

dflash/mask_embedding.pt

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

dflash/model.safetensors

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

dflash/model.safetensors.index.json

@@ -0,0 +1,70 @@
+{
+  "metadata": {
+    "total_size": 2936114304
+  },
+  "weight_map": {
+    "fc.weight": "model.safetensors",
+    "hidden_norm.weight": "model.safetensors",
+    "norm.weight": "model.safetensors",
+    "layers.0.input_layernorm.weight": "model.safetensors",
+    "layers.0.post_attention_layernorm.weight": "model.safetensors",
+    "layers.0.self_attn.q_norm.weight": "model.safetensors",
+    "layers.0.self_attn.k_norm.weight": "model.safetensors",
+    "layers.0.self_attn.q_proj.weight": "model.safetensors",
+    "layers.0.self_attn.k_proj.weight": "model.safetensors",
+    "layers.0.self_attn.v_proj.weight": "model.safetensors",
+    "layers.0.self_attn.o_proj.weight": "model.safetensors",
+    "layers.0.mlp.gate_proj.weight": "model.safetensors",
+    "layers.0.mlp.up_proj.weight": "model.safetensors",
+    "layers.0.mlp.down_proj.weight": "model.safetensors",
+    "layers.0.self_attn.attention_sink_bias": "model.safetensors",
+    "layers.1.input_layernorm.weight": "model.safetensors",
+    "layers.1.post_attention_layernorm.weight": "model.safetensors",
+    "layers.1.self_attn.q_norm.weight": "model.safetensors",
+    "layers.1.self_attn.k_norm.weight": "model.safetensors",
+    "layers.1.self_attn.q_proj.weight": "model.safetensors",
+    "layers.1.self_attn.k_proj.weight": "model.safetensors",
+    "layers.1.self_attn.v_proj.weight": "model.safetensors",
+    "layers.1.self_attn.o_proj.weight": "model.safetensors",
+    "layers.1.mlp.gate_proj.weight": "model.safetensors",
+    "layers.1.mlp.up_proj.weight": "model.safetensors",
+    "layers.1.mlp.down_proj.weight": "model.safetensors",
+    "layers.1.self_attn.attention_sink_bias": "model.safetensors",
+    "layers.2.input_layernorm.weight": "model.safetensors",
+    "layers.2.post_attention_layernorm.weight": "model.safetensors",
+    "layers.2.self_attn.q_norm.weight": "model.safetensors",
+    "layers.2.self_attn.k_norm.weight": "model.safetensors",
+    "layers.2.self_attn.q_proj.weight": "model.safetensors",
+    "layers.2.self_attn.k_proj.weight": "model.safetensors",
+    "layers.2.self_attn.v_proj.weight": "model.safetensors",
+    "layers.2.self_attn.o_proj.weight": "model.safetensors",
+    "layers.2.mlp.gate_proj.weight": "model.safetensors",
+    "layers.2.mlp.up_proj.weight": "model.safetensors",
+    "layers.2.mlp.down_proj.weight": "model.safetensors",
+    "layers.2.self_attn.attention_sink_bias": "model.safetensors",
+    "layers.3.input_layernorm.weight": "model.safetensors",
+    "layers.3.post_attention_layernorm.weight": "model.safetensors",
+    "layers.3.self_attn.q_norm.weight": "model.safetensors",
+    "layers.3.self_attn.k_norm.weight": "model.safetensors",
+    "layers.3.self_attn.q_proj.weight": "model.safetensors",
+    "layers.3.self_attn.k_proj.weight": "model.safetensors",
+    "layers.3.self_attn.v_proj.weight": "model.safetensors",
+    "layers.3.self_attn.o_proj.weight": "model.safetensors",
+    "layers.3.mlp.gate_proj.weight": "model.safetensors",
+    "layers.3.mlp.up_proj.weight": "model.safetensors",
+    "layers.3.mlp.down_proj.weight": "model.safetensors",
+    "layers.3.self_attn.attention_sink_bias": "model.safetensors",
+    "layers.4.input_layernorm.weight": "model.safetensors",
+    "layers.4.post_attention_layernorm.weight": "model.safetensors",
+    "layers.4.self_attn.q_norm.weight": "model.safetensors",
+    "layers.4.self_attn.k_norm.weight": "model.safetensors",
+    "layers.4.self_attn.q_proj.weight": "model.safetensors",
+    "layers.4.self_attn.k_proj.weight": "model.safetensors",
+    "layers.4.self_attn.v_proj.weight": "model.safetensors",
+    "layers.4.self_attn.o_proj.weight": "model.safetensors",
+    "layers.4.mlp.gate_proj.weight": "model.safetensors",
+    "layers.4.mlp.up_proj.weight": "model.safetensors",
+    "layers.4.mlp.down_proj.weight": "model.safetensors",
+    "layers.4.self_attn.attention_sink_bias": "model.safetensors"
+  }
+}

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": false,
+  "eos_token_id": [151643, 151645, 151672],
+  "temperature": 1.0,
+  "top_p": 0.95,
+  "max_new_tokens": 2048,
+  "transformers_version": "4.37.0"
+}

model-00001-of-00036.safetensors

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

model-00002-of-00036.safetensors

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

model-00003-of-00036.safetensors

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

model-00004-of-00036.safetensors

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

model-00005-of-00036.safetensors

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

model-00006-of-00036.safetensors

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

model-00007-of-00036.safetensors

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

model-00008-of-00036.safetensors

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

model-00009-of-00036.safetensors

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

model-00010-of-00036.safetensors

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

model-00011-of-00036.safetensors

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

model-00012-of-00036.safetensors

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

model-00013-of-00036.safetensors

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

model-00014-of-00036.safetensors

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

model-00015-of-00036.safetensors

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

model-00016-of-00036.safetensors

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

model-00017-of-00036.safetensors

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

model-00018-of-00036.safetensors

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

model-00019-of-00036.safetensors

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

model-00020-of-00036.safetensors

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

model-00021-of-00036.safetensors

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

model-00022-of-00036.safetensors

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

model-00023-of-00036.safetensors

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

model-00024-of-00036.safetensors

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

model-00025-of-00036.safetensors

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

model-00026-of-00036.safetensors

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

model-00027-of-00036.safetensors

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

model-00028-of-00036.safetensors

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

model-00029-of-00036.safetensors

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

model-00030-of-00036.safetensors

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

model-00031-of-00036.safetensors

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

model-00032-of-00036.safetensors

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

model-00033-of-00036.safetensors

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