Instructions to use primitive-ai/MiMo-V2.6-Flash-RL-NVFP4 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use primitive-ai/MiMo-V2.6-Flash-RL-NVFP4 with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="primitive-ai/MiMo-V2.6-Flash-RL-NVFP4", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("primitive-ai/MiMo-V2.6-Flash-RL-NVFP4", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use primitive-ai/MiMo-V2.6-Flash-RL-NVFP4 with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "primitive-ai/MiMo-V2.6-Flash-RL-NVFP4"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "primitive-ai/MiMo-V2.6-Flash-RL-NVFP4",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/primitive-ai/MiMo-V2.6-Flash-RL-NVFP4

SGLang

How to use primitive-ai/MiMo-V2.6-Flash-RL-NVFP4 with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "primitive-ai/MiMo-V2.6-Flash-RL-NVFP4" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "primitive-ai/MiMo-V2.6-Flash-RL-NVFP4",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "primitive-ai/MiMo-V2.6-Flash-RL-NVFP4" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "primitive-ai/MiMo-V2.6-Flash-RL-NVFP4",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use primitive-ai/MiMo-V2.6-Flash-RL-NVFP4 with Docker Model Runner:
```
docker model run hf.co/primitive-ai/MiMo-V2.6-Flash-RL-NVFP4
```

jagat-primitive-org commited on 9 days ago

Commit

af9b091

verified ·

1 Parent(s): c83da24

weights: exact MXFP4 -> NVFP4 transcode of XiaomiMiMo/MiMo-V2.6-Flash-RL experts

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

.gitattributes +4 -0
MiMo_V2_6_technical_report.pdf +3 -0
assets/architecture.png +3 -0
audio_tokenizer/chat_template.jinja +120 -0
audio_tokenizer/config.json +68 -0
audio_tokenizer/generation_config.json +9 -0
audio_tokenizer/model.safetensors +3 -0
audio_tokenizer/tokenizer_config.json +267 -0
chat_template.jinja +96 -0
config.json +0 -0
configuration_mimo_v2.py +247 -0
dflash/config.json +56 -0
dflash/dflash.py +379 -0
dflash/dflash_draft_model.safetensors +3 -0
dflash/mask_embedding.pt +3 -0
dflash/model.safetensors.index.json +70 -0
generation_config.json +9 -0
merges.txt +0 -0
model.safetensors.index.json +3 -0
model_mtp.safetensors +3 -0
model_pp0_ep0_shard0.safetensors +3 -0
model_pp0_ep10_shard0.safetensors +3 -0
model_pp0_ep11_shard0.safetensors +3 -0
model_pp0_ep12_shard0.safetensors +3 -0
model_pp0_ep13_shard0.safetensors +3 -0
model_pp0_ep14_shard0.safetensors +3 -0
model_pp0_ep15_shard0.safetensors +3 -0
model_pp0_ep16_shard0.safetensors +3 -0
model_pp0_ep17_shard0.safetensors +3 -0
model_pp0_ep18_shard0.safetensors +3 -0
model_pp0_ep19_shard0.safetensors +3 -0
model_pp0_ep1_shard0.safetensors +3 -0
model_pp0_ep20_shard0.safetensors +3 -0
model_pp0_ep21_shard0.safetensors +3 -0
model_pp0_ep22_shard0.safetensors +3 -0
model_pp0_ep23_shard0.safetensors +3 -0
model_pp0_ep24_shard0.safetensors +3 -0
model_pp0_ep25_shard0.safetensors +3 -0
model_pp0_ep26_shard0.safetensors +3 -0
model_pp0_ep27_shard0.safetensors +3 -0
model_pp0_ep28_shard0.safetensors +3 -0
model_pp0_ep29_shard0.safetensors +3 -0
model_pp0_ep2_shard0.safetensors +3 -0
model_pp0_ep30_shard0.safetensors +3 -0
model_pp0_ep31_shard0.safetensors +3 -0
model_pp0_ep32_shard0.safetensors +3 -0
model_pp0_ep33_shard0.safetensors +3 -0
model_pp0_ep34_shard0.safetensors +3 -0
model_pp0_ep35_shard0.safetensors +3 -0
model_pp0_ep36_shard0.safetensors +3 -0

.gitattributes CHANGED Viewed

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

 *.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
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
+assets/architecture.png filter=lfs diff=lfs merge=lfs -text
+model.safetensors.index.json filter=lfs diff=lfs merge=lfs -text

MiMo_V2_6_technical_report.pdf ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:fb81e6e083801b3358f084ed6be953dc23b0d2e434690f4541d5eae03e01e7af
+size 3046687

assets/architecture.png ADDED Viewed

Git LFS Details

SHA256: d288768e1771fec19b39ed7dbac4adcbbd2e490384d4ad3c58d259b0c6c6bdcc
Pointer size: 131 Bytes
Size of remote file: 405 kB

audio_tokenizer/chat_template.jinja ADDED Viewed

	@@ -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 ADDED Viewed

	@@ -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 ADDED Viewed

	@@ -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 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:077033345d80eef3a315e8d394e0589667e80e4cdaba9bc5a7488410c6657265
+size 1872618384

audio_tokenizer/tokenizer_config.json ADDED Viewed

	@@ -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 ADDED Viewed

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

config.json ADDED Viewed

The diff for this file is too large to render. See raw diff

configuration_mimo_v2.py ADDED Viewed

	@@ -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 ADDED Viewed

	@@ -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 ADDED Viewed

	@@ -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 ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:94d9c02c17e0b469f88699e98424825ebe13a17f07eb3b86727b5e0c8143d287
+size 2936121080

dflash/mask_embedding.pt ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:b35b379fe0497ffdc0d6407f502622237d9346e0d03f457ec574be60f6c2cee0
+size 9882

dflash/model.safetensors.index.json ADDED Viewed

	@@ -0,0 +1,70 @@

+{
+  "metadata": {
+    "total_size": 2936114304
+  },
+  "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 ADDED Viewed

	@@ -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"
+}

merges.txt ADDED Viewed

The diff for this file is too large to render. See raw diff

model.safetensors.index.json ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:54c24ef36b5919dcc02a0cfa751f46dfed56d3a0a223b145a6a67352a6436de5
+size 13979137

model_mtp.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:01bdcd4f6f964e310ef4aceec6db2c16be96b00ad2ecb93a7f0f0effe459ab56
+size 1189405992

model_pp0_ep0_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:e4cb5d053c0ef0059714dca007cdf5c8bc7208ba039439c89598ef82ebff6e5d
+size 13535676488

model_pp0_ep10_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:a1f79a7ffacbf3b25b2c8dcda0421998a9cf5a682276b2b091e4ab8b526d5933
+size 2661557352

model_pp0_ep11_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:797515cfe2149a8f4cc4cec7679a92fabbe2b5d47227f31ff75f2135b7e58430
+size 2661557352

model_pp0_ep12_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:138f0e551937b308b7fc33dc28424757baf85aecba5d2124772ae96e63c9acc6
+size 2661557352

model_pp0_ep13_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:eb1c76e77fac4bd26f12739ffb6603187585c1703a46eb27deca79fa3ce55247
+size 2661557352

model_pp0_ep14_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:4a9ffb6549aa8ce07aa715752091938d14004d55092b00c9e9db9a402095a993
+size 2661557352

model_pp0_ep15_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:f9a1a7318f5339c02aff45836d2fbfe473d89025f8fc1da9127080df90925f04
+size 2661557352

model_pp0_ep16_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:8796b6271d755ca1306cfab3bedb6264f1fd72eef6d0d7dff404b8b4711dda03
+size 2661557352

model_pp0_ep17_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6859251656c31a240bd14c1dffbce11be992cb62862704ba9a3fc907c56174aa
+size 2661557352

model_pp0_ep18_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6729ac6f3cf04f554632abac98408c2139d91cc23fe5fed608aa1e50e203f667
+size 2661557352

model_pp0_ep19_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:eade352e093a12c06c3144e23266a852639be8ab2c6b42a4fd7f9c2b0ee820cb
+size 2661557352

model_pp0_ep1_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:e7c6fbf97ac1e43e9bd9cf6c0d7a49208f1d7d5184b88350ae8b2d43a94ed449
+size 2661555096

model_pp0_ep20_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6adf819c457c21e84d2b197f6fdc3a1a538ebe5110e8902be67b568c52a6cc64
+size 2661557352

model_pp0_ep21_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:5ffe5d134f10e80a55d1df61b26b8586e098aa55c031b30bda826b0fb637006d
+size 2661557352

model_pp0_ep22_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:8c8078ea96fe082b9ca42ba070598072961cfd027e3e79d858c411bea98a49d8
+size 2661557352

model_pp0_ep23_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:0c919297f93f05fff2bfc13839053351b7c09bb86ebef120893ab87424792643
+size 2661557352

model_pp0_ep24_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:a499984ace7caaa91d6a28b9735a6517ed009ae940651b14b8b0836dfa11f291
+size 2661557352

model_pp0_ep25_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:b93968feba089dcf7beff09f35006d90594d7066b90fe27d50853acca40b4bc5
+size 2661559608

model_pp0_ep26_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:de8496820c0fdf7047ee96dc28e8f33954c9bdc13f5fe999370ee7ed17e4f53a
+size 2661559608

model_pp0_ep27_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:e5539cfcc59d302022a4db93b32c77b9930369f776a8062c2e88a365f9b996c3
+size 2661559608

model_pp0_ep28_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:3f9dc0618b425346fc7aac4390c7026343ffd03dc50631935552eef739f99a1f
+size 2661559608

model_pp0_ep29_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:dee7d228c36384a0a7dff029c855634279dcb93a365e999433bde6d1b57b8e6f
+size 2661559608

model_pp0_ep2_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:ed31e61c0d74a907e7b15e9ad4a3b18c0b57edc5478bf2ab63da7bb3ba8ceabb
+size 2661556224

model_pp0_ep30_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:129df6e147809b0e82068ea5e038123d30ff218b27902c5b49a35fae14b57350
+size 2661559608

model_pp0_ep31_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:47d072ca36ae4fb60b4ada54c4b178e4e5df9835e15dc89de51b2270bf8b6393
+size 2661559608

model_pp0_ep32_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6f4514472384bb738d5ef5155d208043902a8b0a3186fe9cb5595634bad5f24b
+size 2661559608

model_pp0_ep33_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:f8c851daf063e14d09b276c214de12e997e5461a6039030b978a6b0430ebe220
+size 2661559608

model_pp0_ep34_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:c8d130b9f112dc21c9af60e9c08fe8ab7de8baecb4d6ae550d69bd236224a311
+size 2661559608

model_pp0_ep35_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:9fc1281c2caed89385b6f43f4cf50137f9dc9da4ac8acd9f3fbbaf0a66c86a54
+size 2661559608

model_pp0_ep36_shard0.safetensors ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:f480ce0d464a5fa3f2971d2f5bf45452d0ed11bb0fd41db84432f49fe8e4adce
+size 2661559608