hermitdave/Agnes-3.0-Flash-MTP-drafter

Agnes-3.0-Flash MTP Drafter

Multi-Token Prediction drafter for speculative decoding with Agnes-3.0-Flash MLX quants.

This is a sidecar drafter — it borrows the target model's embeddings and LM head, so it cannot run standalone.

Architecture

MTP Drafter (~0.85 GB bf16)
├── fc: Linear(2 × hidden_size → hidden_size)     # projects embedding + hidden
├── pre_fc_norm_embedding: RMSNorm
├── pre_fc_norm_hidden: RMSNorm
├── decoder_layer.0 (full attention, 48 Q / 4 KV heads, 256 dim)
│   ├── self_attn: Gated attention (Q/K norms + sigmoid gate)
│   └── mlp: SwiGLU (17408 intermediate)
└── norm: RMSNorm

Forward pass:

1. fc(concat(emb(token), hidden)) → merged representation
2. Self-attention + MLP (borrows target embeddings for token)
3. lm_head(norm(hidden)) → drafted token logits

Usage

With mlx_vlm server

# Install
pip install mlx-vlm

# Start with MTP speculative decoding
python -m mlx_vlm.server \
    --model hermitdave/Agnes-3.0-Flash-MLX-4bit \
    --draft-model hermitdave/Agnes-3.0-Flash-MTP-drafter

With oMLX

⚠️ Not yet supported. oMLX's vlm_mtp_enabled setting only works with vision-language models (VLMs). Agnes is a text-only model, so the MTP drafter cannot be loaded through oMLX settings today. Track oMLX#1089 for generic text MTP support.

Workaround: Use mlx_vlm.server (see above) or the Python API for now.

With Python API

from mlx_lm import load as load_text
from mlx_vlm.speculative.drafters.qwen3_5_mtp.config import Qwen3_5MTPConfig
from mlx_vlm.speculative.drafters.qwen3_5_mtp.qwen3_5_mtp import Qwen3_5MTPDraftModel
import mlx.core as mx, json, safetensors.torch

# Load base
base, tok = load_text("hermitdave/Agnes-3.0-Flash-MLX-4bit")

# Load drafter
cfg = Qwen3_5MTPConfig.from_dict(json.load(open(" Agnes-3.0-Flash-MTP-drafter/config.json")))
mtp = Qwen3_5MTPDraftModel(cfg)
w = safetensors.torch.load_file("Agnes-3.0-Flash-MTP-drafter/model.safetensors")
mtp.load_weights([(k, mx.array(v)) for k, v in w.items()])

# Bind and use
mtp.bind(base)
# ... call mtp.draft_block(...) in speculative decoding loop

Base model compatibility

Base quant	Repo
4-bit	hermitdave/Agnes-3.0-Flash-MLX-4bit
6-bit	hermitdave/Agnes-3.0-Flash-MLX-6bit
8-bit	hermitdave/Agnes-3.0-Flash-MLX-8bit

Performance

Tested on M3 Max 64GB:

• Base model: "The capital of France is" → generates token-by-token
• With MTP drafter: drafts up to 3 tokens per forward pass, then verifies

Expected speedup: 1.5–2.5× on text generation (workload dependent).

Attribution

Drafter weights extracted and converted by Hermes Agent (Nous Research).

Key conversions applied:

• Stripped mtp. prefix from weight names
• Renamed global_attn → self_attn for qwen3_5_mtp compatibility
• Added 1.0 to one-centered RMSNorm weights (HF convention → MLX convention)

Original model: Agnes-AI/Agnes-3.0-Flash (Apache-2.0).