Skip to content

Support Unsloth Gemma (and potentially other) MTP draft models (gemma4-assistant arch) #2229

Description

@markafa

Which version of LM Studio?
Example: LM Studio Version 0.4.20+1 (0.4.20+1)

Which operating system?
MacOS Sequoia 15.7.5

What is the bug?
LM Studio currently cannot properly utilize Unsloth's Gemma MTP draft models (e.g., mtp-gemma-4-E4B-it-Q8_0.gguf). When attaching the MTP file as a draft model, LM Studio treats it as a standard standalone draft LLM rather than passing the --spec-type draft-mtp flag to the llama.cpp backend, causing a KV cache context crash.

Logs
E llama_init_from_model: failed to initialize the context: Gemma4Assistant requires ctx_other to be set (this warning is normal during memory fitting)
E init: the tokens of sequence 3 in the input batch have inconsistent sequence positions:

  • the last position stored in the memory module of the context (i.e. the KV cache) for sequence 3 is X = 7
  • the tokens for sequence 3 in the input batch have a starting position of Y = 0
    E srv decode: failed to process speculative batch

To Reproduce
Steps to reproduce the behavior:
Load gemma-4-E4B-it-qat-UD-Q4_K_XL.gguf (main model).
Attach mtp-gemma-4-E4B-it-Q8_0.gguf (MTP draft model) in the Speculative Decoding / Draft Model setting.
Send a chat message.

Expected Behavior
LM Studio should recognize the gemma4-assistant architecture and automatically pass --spec-type draft-mtp --spec-draft-n-max 4 to the llama.cpp engine, allowing the draft model to share the main model's KV cache.

Working CLI equivalent
This works perfectly in pure llama.cpp using:

./llama-server -m main.gguf --model-draft mtp.gguf --spec-type draft-mtp --spec-draft-n-max 4 -ngl 999

llama.cpp successfully loads and runs these models with the correct flags passed into it while LM Studio possibly does not manage to properly read these models' metadata to see the models actually support MTP loading. This results in MTP option remaining greyed out in the speculative decoding section. I actually tried to make copies of the models and modify their metadata values with python guff tools and load them into LM Studio alongside the original models but nothing worked.

I tried loading the models using the "Draft model" option and change/tweak parameters but this didn't work for apparent reasons as the actual loading parameters should point to --spec-type draft-mtp value, not speculative decoding drafter type as the model structure and behavior is very different for these model architectures.

So when I send a message to models loaded this way it throws the following error though this is expected as the model isn't actually a speculative decoding model but an MTP variant:
Failed to send message
Engine protocol predict stream returned an error: {"code":500,"message":"decode() failed: failed to process speculative batch","type":"server_error"}

LM Studio team could possibly add an advanced settings system which enables the users to pass additional parameters to the underlying inference engines to quickly solve these kind of problems, which could potentially alleviate getting bug reports for many edge cases like this by both allowing us users to keep using LM Studio which we love, and also access the full potential of llama.cpp/mlx engines...

Thank you for you great work!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions