[Model] Add MiMo-V2-Flash support

[Abatom]

Purpose

Add support for MiMo-V2-Flash.

Examples

Example 1

vllm serve XiaomiMiMo/MiMo-V2-Flash \
    --host 0.0.0.0 \
    --port 9001 \
    --seed 1024 \
    --served-model-name base_model \
    --tensor-parallel-size 4 \
    --trust-remote-code \
    --gpu-memory-utilization 0.9 \
    --generation-config vllm

Example 2

vllm serve XiaomiMiMo/MiMo-V2-Flash \
    --host 0.0.0.0 \
    --port 9001 \
    --seed 1024 \
    --served-model-name base_model \
    --data-parallel-size 2 \
    --tensor-parallel-size 4 \
    --enable-expert-parallel \
    --trust-remote-code \
    --gpu-memory-utilization 0.9 \
    --generation-config vllm

Accuracy

GSM8K

lm_eval \
    --model local-chat-completions \
    --tasks gsm8k \
    --num_fewshot 5 \
    --apply_chat_template \
    --model_args model=base_model,base_url=http://0.0.0.0:9001/v1/chat/completions,num_concurrent=100,max_retries=20,tokenized_requests=False,tokenizer_backend=none,max_gen_toks=256

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9128|±  |0.0078|
|     |       |strict-match    |     5|exact_match|↑  |0.9075|±  |0.0080|

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[gemini-code-assist]

Code Review

This pull request introduces support for the MiMo-V2-Flash model. The changes primarily involve adding the model definition in vllm/model_executor/models/mimo_v2_flash.py and making necessary adjustments in vllm/model_executor/layers/linear.py to accommodate features like variable value head sizes and FP8 block shape mismatches. My review has identified a critical issue in the Mixture-of-Experts (MoE) layer detection logic that could lead to incorrect model configurations, and a high-severity issue regarding inconsistent access to the model configuration. Addressing these points will improve the robustness and correctness of the new model implementation.

[jeejeelee]

Some init comments, thank you for contribution

[mergify]

Hi @Abatom, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[jeejeelee]

To add a new LLM model, you also need to :

• Add the model to the documentation
• Register the model in test registry

[mergify]

Hi @Abatom, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @Abatom, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @Abatom, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Documentation preview: https://vllm--30836.org.readthedocs.build/en/30836/

[Abatom]

To add a new LLM model, you also need to :

• Add the model to the documentation
• Register the model in test registry

Done!

[LucasWilkinson]

FYI: #28775 also uses diff q and kv head dims; we can probably support this more naturally in a future follow-up PR

[Isotr0py]

I'm a bit worried that this will cause unexpected behavior for FP8 kernel if we disabled block shape check.

Perhaps we should improve the block shape check for Mimo-V2's edge case instead of just skipping it.

Perhaps @mgoin can give more insights?

[Abatom]

@Isotr0py We tried removing the code above and got the following error.

ValueError: Weight output_partition_size = 192 is not divisible by weight quantization block_n = 128.

[mgoin]

Hmm we support weights that aren't divisible by 128 for other block fp8 models fine, such as kv_a_proj in deepseek, I wonder if it is a specific fused layer

[jeejeelee]

Here are the GSM8K results from my local testing

local-chat-completions (model=XiaomiMiMo/MiMo-V2-Flash,base_url=http://0.0.0.0:8000/v1/chat/completions,num_concurrent=100,max_retries=20,tokenized_requests=False,tokenizer_backend=none,max_gen_toks=256), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9075|±  |0.0080|
|     |       |strict-match    |     5|exact_match|↑  |0.9022|±  |0.0082|

[mergify]

Hi @Abatom, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @Abatom, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[jeejeelee]

Thank you for contribution, let's land this model first and continue improving it in subsequent PRs

[Isotr0py]

Unlock quantization tests CI to see if other models are fine with the temporary hacky fp8 patch.

[mgoin]

I don't think this is valid. When I revert this to apply v = v * self.v_scale before attention in forward pass, I see my gsm8k eval score improve from ~74% to 78%