Instructions to use XiaomiMiMo/MiMo-V2-Flash with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use XiaomiMiMo/MiMo-V2-Flash with Transformers:

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

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

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("XiaomiMiMo/MiMo-V2-Flash", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("XiaomiMiMo/MiMo-V2-Flash", trust_remote_code=True, device_map="auto")
messages = [
    {"role": "user", "content": "Who are you?"},
]
inputs = tokenizer.apply_chat_template(
	messages,
	add_generation_prompt=True,
	tokenize=True,
	return_dict=True,
	return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=40)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))

Inference
Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use XiaomiMiMo/MiMo-V2-Flash with vLLM:

Install from pip and serve model

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

Use Docker

docker model run hf.co/XiaomiMiMo/MiMo-V2-Flash

SGLang

How to use XiaomiMiMo/MiMo-V2-Flash 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 "XiaomiMiMo/MiMo-V2-Flash" \
    --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": "XiaomiMiMo/MiMo-V2-Flash",
		"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 "XiaomiMiMo/MiMo-V2-Flash" \
        --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": "XiaomiMiMo/MiMo-V2-Flash",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use XiaomiMiMo/MiMo-V2-Flash with Docker Model Runner:
```
docker model run hf.co/XiaomiMiMo/MiMo-V2-Flash
```

[Bug] fp8 `weight_scale_inv` shape mismatch for `k_proj` in full-attention layers (TP-sharded block grid)

#44

by Kaixuanliu - opened 10 days ago

Discussion

Kaixuanliu

10 days ago

Summary

In the fp8 checkpoint, the 9 full-attention layers (hybrid_layer_pattern == 0: layers 0, 5, 11, 17, 23, 29, 35, 41, 47) ship a self_attn.k_proj.weight_scale_inv whose block grid does not match the weight it belongs to:

tensor	shape	expected scale grid (`weight_block_size = [128, 128]`)	actual
`model.layers.0.self_attn.k_proj.weight`	`[768, 4096]`	`[6, 32]`	`[8, 32]`

768 = num_key_value_heads(4) * head_dim(192), so ceil(768/128) = 6 block rows are expected, but 8 are stored. All other quantized tensors in the checkpoint are consistent(only these 9 are affected).

Root cause

The scales appear to have been computed per tensor-parallel shard with TP=4 rather than on the full tensor: each rank owns 768 / 4 = 192 rows, which needs ceil(192/128) = 2 block rows, giving 4 * 2 = 8 rows in total. The scale rows therefore pair up per kv head ((0,1), (2,3), (4,5), (6,7)), which is also visible in their magnitudes.

This only shows up for k_proj in full-attention layers because every other projection has a per-shard row count that is a multiple of 128:

q_proj: 12288 / 4 = 3072 -> OK
v_proj: 512 / 4 = 128 -> OK
SWA k_proj: 1536 / 4 = 384 -> OK
full-attn k_proj: 768 / 4 = 192 -> not a multiple of 128

Reproduction

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "XiaomiMiMo/MiMo-V2-Flash"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id,  device_map="auto")
model.generate(**tok("Hello", return_tensors="pt").to(model.device), max_new_tokens=8)

Kaixuanliu

10 days ago

Close it as I suppose it is designed for tp=4, and in this case we cannot cancel the padding.

Upload images, audio, and videos by dragging in the text input, pasting, or clicking here.

Tap or paste here to upload images

· Sign up or log in to comment