Duplicate from XiaomiMiMo/MiMo-V2.6-Pro-RL

Co-authored-by: Bowen Shen <bwshen-mi@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,39 @@
+*.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
+MiMo_V2_6_technical_report.pdf filter=lfs diff=lfs merge=lfs -text
+model.safetensors.index.json 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

MiMo_V2_6_technical_report.pdf

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

README.md

@@ -0,0 +1,237 @@
+---
+license: mit
+language:
+- en
+- zh
+tags:
+- text-generation
+- multimodal
+- vision-language
+- audio
+- agent
+- video-understanding
+- long-context
+- mimo_v2
+- transformers
+library_name: transformers
+---
+
+<br/><br/>
+
+<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="60%" alt="Xiaomi-MiMo" />
+  </picture>
+</div>
+
+<br/>
+
+<div align="center" style="line-height: 1;">
+  |
+  <a href="https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Pro-RL" target="_blank">🤗 HuggingFace</a>
+  &nbsp;|
+  <a href="https://mimo.xiaomi.com/mimo-v2-6" target="_blank">📰 Blog </a>
+  &nbsp;|
+  <a href="https://platform.xiaomimimo.com" target="_blank">🎨 Xiaomi MiMo API Platform </a>
+  &nbsp;|
+  <a href="https://aistudio.xiaomimimo.com" target="_blank">🗨️ Xiaomi MiMo Studio </a>
+  &nbsp;|
+  <a href="https://mimo.xiaomimimo.com/desktop/" target="_blank">💻 Xiaomi MiMo Desktop </a>
+  &nbsp;|
+</div>
+
+<br/>
+
+<div align="center" style="line-height: 1.2;">
+  <strong>Community</strong><br/>
+  <a href="https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro/blob/main/assets/wechat.jpg" target="_blank">WeChat Group</a>
+  &nbsp;|&nbsp;
+  <a href="https://discord.gg/kKC2kNnQEX" target="_blank">Discord</a>
+  &nbsp;|&nbsp;
+  <a href="https://t.me/+3T-I0pekOVIyNDBl" target="_blank">Telegram</a>
+  &nbsp;|&nbsp;
+  <a href="https://www.reddit.com/r/XiaomiMiMo_Official/" target="_blank">Reddit</a>
+</div>
+
+<br/>
+
+# MiMo-V2.6-Pro-RL
+
+**Scaling Reinforcement Learning Toward Self-Improvement**
+
+<p align="center">
+  <a href="https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Pro-RL/blob/main/MiMo_V2_6_technical_report.pdf"><b>Technical Report</b></a>
+</p>
+
+## 1. Introduction
+
+MiMo-V2.6-Pro-RL is the flagship checkpoint of the MiMo-V2.6 series. The series is built to **scale reinforcement learning toward self-improvement** — scaling RL compute, environment diversity, and grader compute together, so the model keeps expanding its capability frontier through exploration and feedback. Key features include:
+
+- **Native Omnimodal + Long Horizon**: Text, image, video, and audio in one model; 1M tokens for long repositories, tool traces, and multi-session agent runs.
+- **You Only RL Once**: One mixed RL run across coding, general agents, visual, and cybersecurity — not separate per-domain runs. Tasks and multiple harnesses are mixed in the same batch so capabilities reinforce each other and strategies transfer to harnesses never seen in training.
+- **Scaling RL Compute**: Fully asynchronous Group Relative Policy Optimization (GRPO) on very large batches — 1,568 prompts × 16 rollouts per step, billions of tokens per update.
+- **Groupwise Agentic Grading (Self-Improvement Loop)**: Binary pass/fail cannot rank passing solutions, so the reward signal itself is scaled. An agentic grader compares rollouts *within each group*: **Groupwise Reward Synthesis (GRS)** builds task-specific rubrics offline from contrasting rollouts and fuses rubric quality with test outcomes; **Groupwise Advantage Redistribution (GAR)** ranks passing trajectories online and moves advantage toward higher-quality solutions. Judged against the policy’s own samples, this closes a self-improvement loop and steers toward shorter paths and fewer tokens per task.
+- **Aligned RL**: Cold start from self-correction — the model reflects on and rewrites its own misaligned turns into grounded next steps. Throughout RL, environment hardening, adversarial screening, and verifier cross-checks keep the loop honest against reward hacking.
+- **Multi-Prefix Multi-Teacher On-Policy Distillation (MOPD2)**: After mixed RL, MOPD2 combines autonomous student rollouts with prefix-conditioned single-turn rollouts (Teacher-Prefix and SFT-Prefix), reusing histories from teacher trajectories and SFT demonstrations so decision points train without regenerating preceding turns — extending capabilities to hard-to-verify tasks.
+
+## Model Summary
+
+- **Architecture**: Sparse MoE (Mixture of Experts), 1.02T total / 42B activated parameters
+- **Context Length**: 1M tokens
+- **Modalities**: Text, Image, Video, Audio
+- **Vision Encoder**: 681M-param MiMo ViT (28 layers: 24 SWA + 4 Full)
+- **Audio Encoder**: 308M AudioTokenizer + 127M audio patch encoder
+- **Multi-Token Prediction (MTP)**: 5-layer speculative decoder
+
+![Figure 1: MiMo-V2.6 architecture — omni encoders, hybrid SWA backbone, and MTP blocks](assets/architecture.png)
+
+*Figure 1. MiMo-V2.6 architecture.*
+
+## 2. Downloads
+
+| Model | Download |
+| --- | --- |
+| **MiMo-V2.6-Pro-RL** | [🤗 HuggingFace](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Pro-RL) · [🤖 ModelScope](https://www.modelscope.cn/models/XiaomiMiMo/MiMo-V2.6-Pro-RL) |
+| **MiMo-V2.6-Flash-RL** | [🤗 HuggingFace](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL) · [🤖 ModelScope](https://www.modelscope.cn/models/XiaomiMiMo/MiMo-V2.6-Flash-RL) |
+
+## 3. Evaluation Results
+
+| Benchmark | MiMo-V2.6 Pro | MiMo-V2.6 Flash | MiMo-V2.5 Pro | Claude Opus 5 | GPT-5.6 Sol | Claude Fable 5 |
+| --- | --- | --- | --- | --- | --- | --- |
+| **Code Agent** | | | | | | |
+| DeepSWE v1.1 | 71.9 | 67.9 | 19.0 | 74.0 | 73.0 | 70.0 |
+| ProgramBench | 26.5 | 26.0 | 12.5 | 37.0 | 25.0 | 33.0 |
+| MiMo Code Bench | 63.2 | 61.2 | 40.4 | 68.6 | 59.3 | - |
+| **General Agent** | | | | | | |
+| AutomationBench v1.0.6 | 53.1 | 52.3 | 16.0 | 50.3 | 45.8 | 46.2 |
+| Toolathlon-Verified | 76.9 | 73.6 | 49.1 | 80.6 | 74.9 | 77.9 |
+| GDPval-AA 2.1 | 1673 | - | 1107 | 1708 | 1588 | 1595 |
+| Agents’ Last Exam | 31.6 | 27.6 | 13.2 | 31.6 | 30.8 | 25.7 |
+| Terminal Bench 4.0 | 34.9 | 28.8 | 1.5 | 49.0 | 39.9 | 42.4 |
+| Terminal Bench 2.1 | 89.9 | 87.6 | 65.2 | 89.1 | 88.8 | 84.3 |
+| OSWorld-Verified | 82.0 | 80.8 | - | 83.4 | 83.0 | 86.0 |
+| JobBench | 62.0 | 61.2 | 25.0 | 65.7 | 45.4 | 57.4 |
+| **Cybersecurity** | | | | | | |
+| CyberGym | 94.0 | 95.1 | 40.0 | - | - | - |
+| MiMo Cyber Bench | 80.2 | 77.2 | 0.0 | - | - | - |
+| ExploitGym | 17.8 | 6.0 | 0.2 | 22.1 | 30.3 | 28.4 |
+| ExploitBench | 47.9 | 25.3 | 16.6 | 70.0 | 78.5 | 78.0 |
+| SEC Bench Pro | 66.3 | 47.5 | 17.7 | - | 79.1 | - |
+| **Visual Agent** | | | | | | |
+| MiMo VisualCoding | 72.3 | 71.5 | - | 70.0 | 73.4 | 69.1 |
+
+## 4. Model Architecture
+
+### LLM Backbone
+
+| Component | MiMo-V2.6-Pro-RL |
+| --- | --- |
+| Layers (Total / SWA / GA) | 70 / 60 / 10 |
+| Hidden Size | 6144 |
+| SWA Heads (Q/KV) | 128 / 8 |
+| GA Heads (Q/KV) | 128 / 8 |
+| Head Dimensions (QK / V) | 192 / 128 |
+| Sliding Window Size | 128 |
+| Routed Experts (Total / Activated) | 384 / 8 |
+| Max Context Length | 1M |
+| MTP / Speculative Decoder | 5 SWA layers, window 1024 |
+
+The first Transformer block uses global attention with a dense FFN. Remaining blocks interleave local SWA and GA; both use sparse MoE FFNs without shared experts.
+
+### Vision Encoder (MiMo ViT)
+
+| Configuration | Value |
+| --- | --- |
+| Layers (Total / SWA / GA) | 28 / 24 / 4 |
+| Hidden Size | 1280 |
+| Attention Heads (Q / KV) | 32 / 8 |
+| Head Dimension | 64 |
+| Patch Size (T × H × W) | 2 × 16 × 16 |
+| Sliding Window (Left / Right) | 64 / 64 |
+| Spatial Merge Size | 2 × 2 |
+| Parameters | 681M |
+
+### Audio Encoders
+
+AudioTokenizer encoder: 24 layers (12 SWA / 12 GA), hidden 1024, 20 RVQ codebooks, 308M parameters. Audio patch encoder: 6 layers, 127M parameters; four frames per patch (25 Hz → 6.25 Hz).
+
+### Speculative Decoder
+
+5-layer SWA MTP drafter (DFlash-style). Predicts 7 subsequent tokens per forward pass for parallel verification.
+
+## 5. Deployment
+
+For best performance, follow the [SGLang MiMo cookbook](https://docs.sglang.io/cookbook/autoregressive/Xiaomi/MiMo-V2.5). Docker image: `lmsysorg/sglang:latest`.
+
+### SGLang
+
+```bash
+sglang serve \
+  --trust-remote-code \
+  --model-path XiaomiMiMo/MiMo-V2.6-Pro-RL \
+  --tp 16 \
+  --dp 2 \
+  --enable-dp-attention \
+  --mm-enable-dp-encoder \
+  --ep 16 \
+  --moe-a2a-backend deepep \
+  --moe-dense-tp-size 1 \
+  --mem-fraction-static 0.7 \
+  --max-running-requests 128 \
+  --chunked-prefill-size 32768 \
+  --page-size 64 \
+  --swa-full-tokens-ratio 0.3 \
+  --speculative-algorithm EAGLE \
+  --speculative-num-steps 3 \
+  --speculative-eagle-topk 1 \
+  --speculative-num-draft-tokens 4 \
+  --enable-multi-layer-eagle \
+  --reasoning-parser mimo \
+  --tool-call-parser mimo \
+  --host 0.0.0.0 \
+  --port 30000 \
+  --nnodes 2 \
+  --node-rank <node-rank> \
+  --dist-init-addr <node0-ip>:20000
+```
+
+### vLLM
+
+Follow the [vLLM MiMo-V2.5 recipe](https://recipes.vllm.ai/XiaomiMiMo/MiMo-V2.5). Pre-built image: `docker pull vllm/vllm-openai:mimov25-cu129`.
+
+```bash
+vllm serve XiaomiMiMo/MiMo-V2.6-Pro-RL \
+  --tensor-parallel-size 8 \
+  --trust-remote-code \
+  --gpu-memory-utilization 0.95 \
+  --max-model-len auto \
+  --reasoning-parser mimo \
+  --tool-call-parser mimo \
+  --enable-auto-tool-choice \
+  --generation-config vllm
+```
+
+Recommended sampling: `temperature=1.0`, `top_p=0.95`.
+
+Also available in AI Studio, MiMo Code, Xiaomi MiMo Desktop, Xiaomi MiMo Open Platform API, and OpenRouter.
+
+## Citation
+
+```bibtex
+@misc{mimo2026v26pro,
+  title={MiMo-V2.6-Pro-RL},
+  author={{Xiaomi MiMo Team}},
+  year={2026},
+  howpublished={\url{https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Pro-RL}},
+}
+```
+
+## Contact
+
+For questions or feedback, reach us at [mimo@xiaomi.com](mailto:mimo@xiaomi.com) or join our community:
+
+- [WeChat Group](https://work.weixin.qq.com/apph5/external_room/join/group_mng?plg_id=c417f99bd9014b5dd894daa8bfe19790&)
+- [Discord](https://discord.gg/WX2R2uNp)
+- [Telegram](https://t.me/+3T-I0pekOVIyNDBl)
+- [Reddit](https://www.reddit.com/r/XiaomiMiMo_Official/)

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,95 @@
+{%- 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 -%}

config.json

@@ -0,0 +1,441 @@
+{
+  "add_full_attention_sink_bias": false,
+  "add_swa_attention_sink_bias": true,
+  "architectures": [
+    "MiMoV2ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_chunk_size": 128,
+  "attention_dropout": 0.0,
+  "attention_projection_layout": "fused_qkv",
+  "attention_value_scale": 0.612,
+  "audio_config": {
+    "add_post_norm": true,
+    "audio_channels": 20,
+    "audio_segment_size": 6000,
+    "group_size": 4,
+    "input_full_attention": true,
+    "input_local_attn_heads": 16,
+    "input_local_dim": 1024,
+    "input_local_head_dim": 64,
+    "input_local_hidden_dropout": 0.0,
+    "input_local_intermediate_size": 4096,
+    "input_local_layers": 6,
+    "out_hidden_size": 6144,
+    "partial_rotary_factor": 1.0,
+    "projection_layers": 2,
+    "rope_theta": 640000,
+    "speech_vocab_size": "1280",
+    "speech_zeroemb_idx": "1024"
+  },
+  "audio_end_token_id": 151674,
+  "audio_start_token_id": 151673,
+  "audio_token_id": 151669,
+  "auto_map": {
+    "AutoConfig": "configuration_mimo_v2.MiMoV2Config",
+    "AutoModel": "modeling_mimo_v2.MiMoV2Model",
+    "AutoModelForCausalLM": "modeling_mimo_v2.MiMoV2ForCausalLM"
+  },
+  "moe_router_dtype": "bfloat16",
+  "bos_token_id": null,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "head_dim": 192,
+  "hidden_act": "silu",
+  "hidden_size": 6144,
+  "hybrid_block_size": null,
+  "hybrid_layer_pattern": [
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0
+  ],
+  "image_token_id": 151655,
+  "initializer_range": 0.02,
+  "intermediate_size": 16384,
+  "layernorm_epsilon": 1e-05,
+  "max_position_embeddings": 1048576,
+  "model_type": "mimo_v2",
+  "moe_intermediate_size": 2048,
+  "moe_layer_freq": [
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1
+  ],
+  "n_group": 1,
+  "n_routed_experts": 384,
+  "n_shared_experts": null,
+  "norm_topk_prob": true,
+  "num_attention_heads": 128,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 70,
+  "num_key_value_heads": 8,
+  "pad_token_id": 151643,
+  "partial_rotary_factor": 0.334,
+  "processor_config": {
+    "audio_avg_pooler": 2,
+    "audio_channels": 20,
+    "audio_end_token_id": 151674,
+    "audio_fmax": null,
+    "audio_fmin": 0,
+    "audio_group_size": 4,
+    "audio_hop_length": 240,
+    "audio_input_id_per_second": 25.0,
+    "audio_kernel_size": 3,
+    "audio_n_mels": 128,
+    "audio_nfft": 960,
+    "audio_sampling_rate": 24000,
+    "audio_segment_size": 6000,
+    "audio_start_token_id": 151673,
+    "audio_stride_size": 2,
+    "audio_token_id": 151669,
+    "audio_window_size": 960,
+    "audio_zeroemb_idx": [
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024,
+      1024
+    ],
+    "fps": 1.0,
+    "image_max_pixels": 8388608,
+    "image_min_pixels": 8192,
+    "image_token_id": 151655,
+    "max_frames": 3600,
+    "merge_size": 2,
+    "min_frames": null,
+    "num_frames": null,
+    "pad_token_id": 151643,
+    "patch_size": 16,
+    "rope_type": "rope",
+    "temporal_compression_ratio": 1,
+    "temporal_patch_size": 2,
+    "use_per_grid_t_timestamps": false,
+    "use_video_timestamps": true,
+    "video_audio_interleave_length": 0.0,
+    "video_end_token_id": 151671,
+    "video_max_pixels": 8388608,
+    "video_min_pixels": 8192,
+    "video_process_num_threads": 16,
+    "video_start_token_id": 151670,
+    "video_token_id": 151656,
+    "video_tokens_per_second": 2,
+    "video_total_max_pixels": 268435456,
+    "vision_end_token_id": 151653,
+    "vision_start_token_id": 151652
+  },
+  "quantization_config": {
+    "activation_scheme": "dynamic",
+    "fmt": "e4m3",
+    "ignored_layers": [
+      "model.layers.0.self_attn.o_proj",
+      "model.layers.1.self_attn.o_proj",
+      "model.layers.10.self_attn.o_proj",
+      "model.layers.11.self_attn.o_proj",
+      "model.layers.12.self_attn.o_proj",
+      "model.layers.13.self_attn.o_proj",
+      "model.layers.14.self_attn.o_proj",
+      "model.layers.15.self_attn.o_proj",
+      "model.layers.16.self_attn.o_proj",
+      "model.layers.17.self_attn.o_proj",
+      "model.layers.18.self_attn.o_proj",
+      "model.layers.19.self_attn.o_proj",
+      "model.layers.2.self_attn.o_proj",
+      "model.layers.20.self_attn.o_proj",
+      "model.layers.21.self_attn.o_proj",
+      "model.layers.22.self_attn.o_proj",
+      "model.layers.23.self_attn.o_proj",
+      "model.layers.24.self_attn.o_proj",
+      "model.layers.25.self_attn.o_proj",
+      "model.layers.26.self_attn.o_proj",
+      "model.layers.27.self_attn.o_proj",
+      "model.layers.28.self_attn.o_proj",
+      "model.layers.29.self_attn.o_proj",
+      "model.layers.3.self_attn.o_proj",
+      "model.layers.30.self_attn.o_proj",
+      "model.layers.31.self_attn.o_proj",
+      "model.layers.32.self_attn.o_proj",
+      "model.layers.33.self_attn.o_proj",
+      "model.layers.34.self_attn.o_proj",
+      "model.layers.35.self_attn.o_proj",
+      "model.layers.36.self_attn.o_proj",
+      "model.layers.37.self_attn.o_proj",
+      "model.layers.38.self_attn.o_proj",
+      "model.layers.39.self_attn.o_proj",
+      "model.layers.4.self_attn.o_proj",
+      "model.layers.40.self_attn.o_proj",
+      "model.layers.41.self_attn.o_proj",
+      "model.layers.42.self_attn.o_proj",
+      "model.layers.43.self_attn.o_proj",
+      "model.layers.44.self_attn.o_proj",
+      "model.layers.45.self_attn.o_proj",
+      "model.layers.46.self_attn.o_proj",
+      "model.layers.47.self_attn.o_proj",
+      "model.layers.48.self_attn.o_proj",
+      "model.layers.49.self_attn.o_proj",
+      "model.layers.5.self_attn.o_proj",
+      "model.layers.50.self_attn.o_proj",
+      "model.layers.51.self_attn.o_proj",
+      "model.layers.52.self_attn.o_proj",
+      "model.layers.53.self_attn.o_proj",
+      "model.layers.54.self_attn.o_proj",
+      "model.layers.55.self_attn.o_proj",
+      "model.layers.56.self_attn.o_proj",
+      "model.layers.57.self_attn.o_proj",
+      "model.layers.58.self_attn.o_proj",
+      "model.layers.59.self_attn.o_proj",
+      "model.layers.6.self_attn.o_proj",
+      "model.layers.60.self_attn.o_proj",
+      "model.layers.61.self_attn.o_proj",
+      "model.layers.62.self_attn.o_proj",
+      "model.layers.63.self_attn.o_proj",
+      "model.layers.64.self_attn.o_proj",
+      "model.layers.65.self_attn.o_proj",
+      "model.layers.66.self_attn.o_proj",
+      "model.layers.67.self_attn.o_proj",
+      "model.layers.68.self_attn.o_proj",
+      "model.layers.69.self_attn.o_proj",
+      "model.layers.7.self_attn.o_proj",
+      "model.layers.8.self_attn.o_proj",
+      "model.layers.9.self_attn.o_proj",
+      "model.decoder.self_attn.o_proj"
+    ],
+    "mxfp4_block_size": 32,
+    "quant_method": "fp8",
+    "store_dtype": "mxfp4",
+    "weight_block_size": [
+      128,
+      128
+    ]
+  },
+  "rope_parameters": {
+    "partial_rotary_factor": 0.334,
+    "rope_theta": 10000000,
+    "rope_type": "default",
+    "type": "default"
+  },
+  "rope_theta": 10000000,
+  "routed_scaling_factor": null,
+  "scoring_func": "sigmoid",
+  "sliding_window": 128,
+  "sliding_window_size": 128,
+  "swa_head_dim": 192,
+  "swa_num_attention_heads": 128,
+  "swa_num_key_value_heads": 8,
+  "swa_rope_theta": 10000,
+  "swa_v_head_dim": 128,
+  "tie_word_embeddings": false,
+  "topk_group": 1,
+  "topk_method": "noaux_tc",
+  "transformers_version": "5.3.0",
+  "use_cache": true,
+  "v_head_dim": 128,
+  "video_token_id": 151656,
+  "vision_config": {
+    "depth": 28,
+    "fullatt_block_indexes": [
+      0,
+      9,
+      18,
+      27
+    ],
+    "hidden_act": "silu",
+    "hidden_size": 1280,
+    "in_chans": 3,
+    "intermediate_size": 4608,
+    "num_heads": 32,
+    "num_key_value_heads": 8,
+    "num_query_groups": 4,
+    "out_hidden_size": 6144,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "spatial_patch_size": 16,
+    "temporal_patch_size": 2,
+    "tokens_per_second": 2,
+    "use_sink": true,
+    "visual_token_window_size": 64,
+    "vit_window_attn_types": [
+      -1,
+      0,
+      0,
+      0,
+      0,
+      1,
+      1,
+      1,
+      1,
+      -1,
+      0,
+      0,
+      0,
+      0,
+      1,
+      1,
+      1,
+      1,
+      -1,
+      0,
+      0,
+      0,
+      0,
+      1,
+      1,
+      1,
+      1,
+      -1
+    ],
+    "window_size": 128
+  },
+  "vision_end_token_id": 151653,
+  "vision_model_type": "mimovl",
+  "vision_start_token_id": 151652,
+  "vocab_size": 152576
+}

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": 6144,
+  "intermediate_size": 16384,
+  "num_hidden_layers": 5,
+  "num_attention_heads": 128,
+  "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,
+      15,
+      31,
+      47,
+      69
+    ],
+    "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": 70,
+  "target_hidden_size": 6144,
+  "vocab_size": 152576,
+  "max_position_embeddings": 1048576,
+  "rope_theta": 10000.0,
+  "rms_norm_eps": 1e-05,
+  "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/dflash_draft_model.safetensors

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

dflash/mask_embedding.pt

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

dflash/model.safetensors.index.json

@@ -0,0 +1,70 @@
+{
+  "metadata": {
+    "total_size": 5536632576
+  },
+  "weight_map": {
+    "fc.weight": "dflash_draft_model.safetensors",
+    "hidden_norm.weight": "dflash_draft_model.safetensors",
+    "norm.weight": "dflash_draft_model.safetensors",
+    "layers.0.input_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.0.post_attention_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.0.self_attn.q_norm.weight": "dflash_draft_model.safetensors",
+    "layers.0.self_attn.k_norm.weight": "dflash_draft_model.safetensors",
+    "layers.0.self_attn.q_proj.weight": "dflash_draft_model.safetensors",
+    "layers.0.self_attn.k_proj.weight": "dflash_draft_model.safetensors",
+    "layers.0.self_attn.v_proj.weight": "dflash_draft_model.safetensors",
+    "layers.0.self_attn.o_proj.weight": "dflash_draft_model.safetensors",
+    "layers.0.mlp.gate_proj.weight": "dflash_draft_model.safetensors",
+    "layers.0.mlp.up_proj.weight": "dflash_draft_model.safetensors",
+    "layers.0.mlp.down_proj.weight": "dflash_draft_model.safetensors",
+    "layers.0.self_attn.attention_sink_bias": "dflash_draft_model.safetensors",
+    "layers.1.input_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.1.post_attention_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.1.self_attn.q_norm.weight": "dflash_draft_model.safetensors",
+    "layers.1.self_attn.k_norm.weight": "dflash_draft_model.safetensors",
+    "layers.1.self_attn.q_proj.weight": "dflash_draft_model.safetensors",
+    "layers.1.self_attn.k_proj.weight": "dflash_draft_model.safetensors",
+    "layers.1.self_attn.v_proj.weight": "dflash_draft_model.safetensors",
+    "layers.1.self_attn.o_proj.weight": "dflash_draft_model.safetensors",
+    "layers.1.mlp.gate_proj.weight": "dflash_draft_model.safetensors",
+    "layers.1.mlp.up_proj.weight": "dflash_draft_model.safetensors",
+    "layers.1.mlp.down_proj.weight": "dflash_draft_model.safetensors",
+    "layers.1.self_attn.attention_sink_bias": "dflash_draft_model.safetensors",
+    "layers.2.input_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.2.post_attention_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.2.self_attn.q_norm.weight": "dflash_draft_model.safetensors",
+    "layers.2.self_attn.k_norm.weight": "dflash_draft_model.safetensors",
+    "layers.2.self_attn.q_proj.weight": "dflash_draft_model.safetensors",
+    "layers.2.self_attn.k_proj.weight": "dflash_draft_model.safetensors",
+    "layers.2.self_attn.v_proj.weight": "dflash_draft_model.safetensors",
+    "layers.2.self_attn.o_proj.weight": "dflash_draft_model.safetensors",
+    "layers.2.mlp.gate_proj.weight": "dflash_draft_model.safetensors",
+    "layers.2.mlp.up_proj.weight": "dflash_draft_model.safetensors",
+    "layers.2.mlp.down_proj.weight": "dflash_draft_model.safetensors",
+    "layers.2.self_attn.attention_sink_bias": "dflash_draft_model.safetensors",
+    "layers.3.input_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.3.post_attention_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.3.self_attn.q_norm.weight": "dflash_draft_model.safetensors",
+    "layers.3.self_attn.k_norm.weight": "dflash_draft_model.safetensors",
+    "layers.3.self_attn.q_proj.weight": "dflash_draft_model.safetensors",
+    "layers.3.self_attn.k_proj.weight": "dflash_draft_model.safetensors",
+    "layers.3.self_attn.v_proj.weight": "dflash_draft_model.safetensors",
+    "layers.3.self_attn.o_proj.weight": "dflash_draft_model.safetensors",
+    "layers.3.mlp.gate_proj.weight": "dflash_draft_model.safetensors",
+    "layers.3.mlp.up_proj.weight": "dflash_draft_model.safetensors",
+    "layers.3.mlp.down_proj.weight": "dflash_draft_model.safetensors",
+    "layers.3.self_attn.attention_sink_bias": "dflash_draft_model.safetensors",
+    "layers.4.input_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.4.post_attention_layernorm.weight": "dflash_draft_model.safetensors",
+    "layers.4.self_attn.q_norm.weight": "dflash_draft_model.safetensors",
+    "layers.4.self_attn.k_norm.weight": "dflash_draft_model.safetensors",
+    "layers.4.self_attn.q_proj.weight": "dflash_draft_model.safetensors",
+    "layers.4.self_attn.k_proj.weight": "dflash_draft_model.safetensors",
+    "layers.4.self_attn.v_proj.weight": "dflash_draft_model.safetensors",
+    "layers.4.self_attn.o_proj.weight": "dflash_draft_model.safetensors",
+    "layers.4.mlp.gate_proj.weight": "dflash_draft_model.safetensors",
+    "layers.4.mlp.up_proj.weight": "dflash_draft_model.safetensors",
+    "layers.4.mlp.down_proj.weight": "dflash_draft_model.safetensors",
+    "layers.4.self_attn.attention_sink_bias": "dflash_draft_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.safetensors.index.json

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

model_mtp.safetensors

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

model_pp0_ep0_shard0.safetensors

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

model_pp0_ep0_shard1.safetensors

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

model_pp0_ep100_shard0.safetensors

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

model_pp0_ep101_shard0.safetensors

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

model_pp0_ep102_shard0.safetensors

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

model_pp0_ep103_shard0.safetensors

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

model_pp0_ep104_shard0.safetensors

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

model_pp0_ep105_shard0.safetensors

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

model_pp0_ep106_shard0.safetensors

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

model_pp0_ep107_shard0.safetensors

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

model_pp0_ep108_shard0.safetensors

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

model_pp0_ep109_shard0.safetensors

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

model_pp0_ep10_shard0.safetensors

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

model_pp0_ep110_shard0.safetensors

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

model_pp0_ep111_shard0.safetensors

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

model_pp0_ep112_shard0.safetensors

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

model_pp0_ep113_shard0.safetensors

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

model_pp0_ep114_shard0.safetensors

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

model_pp0_ep115_shard0.safetensors

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

model_pp0_ep116_shard0.safetensors

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

model_pp0_ep117_shard0.safetensors

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

model_pp0_ep118_shard0.safetensors

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

model_pp0_ep119_shard0.safetensors

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

model_pp0_ep11_shard0.safetensors

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

model_pp0_ep120_shard0.safetensors

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

model_pp0_ep121_shard0.safetensors

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

model_pp0_ep122_shard0.safetensors

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

model_pp0_ep123_shard0.safetensors

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

model_pp0_ep124_shard0.safetensors

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