use anyhow::{bail, ensure, Context, Result};
use mlx_native::gguf::GgufFile;
use mlx_native::{MlxBuffer, MlxDevice};
use super::ffn::{DenseFfnWeights, MoeFfnShape};
use super::gpu_ffn::{DenseFfnWeightsGpu, DenseFfnWeightsGpuQ, MoeFfnWeightsGpuQ};
use super::gpu_full_attn::{upload_f32_weight, FullAttnQGateWeightsGpu};
use super::mtp::{MtpFfnWeightsGpu, MtpFullAttnWeightsGpu, MtpQGateWeightsGpu, MtpWeights};
use super::weight_loader::{
dense_ffn_storage, dense_ffn_tensor_types, load_dense_ffn_quantized, load_f32_tensor,
load_moe_ffn_quantized, load_native_projection, qwen35_moe_expert_type_supported,
validate_native_projection_info, DenseFfnStorage,
};
use super::Qwen35Config;
use crate::serve::forward_mlx_shared::MlxQWeight;
pub fn load_mtp_weights_if_present(
gguf: &GgufFile,
cfg: &Qwen35Config,
device: &MlxDevice,
) -> Result<Option<MtpWeights>> {
load_mtp_weights_if_present_with_shared_head(gguf, cfg, device, None)
}
pub fn load_mtp_weights_if_present_with_shared_head(
gguf: &GgufFile,
cfg: &Qwen35Config,
device: &MlxDevice,
main_output_head: Option<&MlxQWeight>,
) -> Result<Option<MtpWeights>> {
if cfg.mtp_num_hidden_layers == 0 {
return Ok(None);
}
validate_mtp_tensor_topology(gguf, cfg)?;
let layer_index = cfg.num_hidden_layers;
let loaded_tensor_names = mtp_tensor_names(gguf, layer_index);
let h = cfg.hidden_size as usize;
let p = format!("blk.{layer_index}");
let nextn = format!("{p}.nextn");
let enorm = load_norm_gpu(gguf, &format!("{nextn}.enorm.weight"), h, device)?;
let hnorm = load_norm_gpu(gguf, &format!("{nextn}.hnorm.weight"), h, device)?;
let embed_tokens_tname = format!("{nextn}.embed_tokens.weight");
let embed_tokens = if cfg.mtp_use_dedicated_embeddings {
let info = gguf.tensor_info(&embed_tokens_tname).ok_or_else(|| {
anyhow::anyhow!("MTP dedicated embedding tensor `{embed_tokens_tname}` is missing")
})?;
ensure!(
info.shape.len() == 2 && info.shape[1] == h,
"{embed_tokens_tname} shape {:?} is not [vocab,{h}]",
info.shape
);
let buf =
crate::serve::forward_mlx_shared::load_gguf_qweight(gguf, &embed_tokens_tname, device)
.with_context(|| {
format!(
"MTP loader expected dedicated `{embed_tokens_tname}` because \
mtp_use_dedicated_embeddings=True"
)
})?;
super::forward_gpu::ensure_native_embedding_admitted(&buf)
.with_context(|| format!("admit direct execution for {embed_tokens_tname}"))?;
super::weight_pool::register_weight_buffer(device, &buf.buffer)
.with_context(|| format!("register {embed_tokens_tname}"))?;
tracing::info!(
mtp_layer = layer_index,
mtp_use_dedicated_embeddings = true,
tensor = %embed_tokens_tname,
"qwen35 MTP loader: dedicated embed_tokens"
);
Some(buf)
} else {
if gguf.tensor_info(&embed_tokens_tname).is_some() {
bail!(
"qwen35 MTP loader: mtp_use_dedicated_embeddings=False but `{embed_tokens_tname}` \
is present in the GGUF — convert pipeline is inconsistent. Re-emit without \
the dedicated tensor or set the metadata key to true."
);
}
tracing::info!(
mtp_layer = layer_index,
mtp_use_dedicated_embeddings = false,
"qwen35 MTP loader: sharing main token_embd (no dedicated nextn.embed_tokens)"
);
None
};
let (eh_proj, eh_proj_ggml_type) =
load_native_projection(gguf, &format!("{nextn}.eh_proj.weight"), h, 2 * h, device)?;
let shared_head_norm =
load_norm_gpu(gguf, &format!("{nextn}.shared_head_norm.weight"), h, device)?;
let shared_head_head_tname = format!("{nextn}.shared_head_head.weight");
let (shared_head_head, shared_head_head_ggml_type, vocab_size, shared_head_head_source) =
if let Some(info) = gguf.tensor_info(&shared_head_head_tname) {
ensure!(
info.shape.len() == 2 && info.shape[1] == h,
"{shared_head_head_tname} shape {:?} is not [vocab, {h}]",
info.shape
);
let vocab = info.shape[0];
let (buffer, ggml_type) =
load_native_projection(gguf, &shared_head_head_tname, vocab, h, device)?;
(
buffer,
ggml_type,
vocab as u32,
shared_head_head_tname.clone(),
)
} else if !cfg.mtp_use_dedicated_embeddings {
let main_lm = if gguf.tensor_info("output.weight").is_some() {
"output.weight"
} else {
"token_embd.weight"
};
if let Some(main) = main_output_head {
ensure!(
main.affine.is_none() && main.info.cols == h,
"qwen35 MTP loader: supplied shared output head is not a native [vocab,{h}] GGUF projection"
);
(
main.buffer.clone(),
main.info.ggml_dtype,
main.info.rows as u32,
main_lm.to_string(),
)
} else {
let info = gguf.tensor_info(main_lm).ok_or_else(|| {
anyhow::anyhow!(
"qwen35 MTP loader: shared head missing and resolved main head {main_lm} absent"
)
})?;
ensure!(
info.shape.len() == 2 && info.shape[1] == h,
"{main_lm} shape {:?} is not [vocab, {h}]",
info.shape
);
let vocab = info.shape[0];
let (buffer, ggml_type) = load_native_projection(gguf, main_lm, vocab, h, device)?;
(buffer, ggml_type, vocab as u32, main_lm.to_string())
}
} else {
bail!(
"qwen35 MTP loader: `{shared_head_head_tname}` is missing AND \
mtp_use_dedicated_embeddings=True — cannot resolve the MTP final \
projection. Re-emit the GGUF with the dedicated tensor or set the \
flag to false."
);
};
tracing::info!(
mtp_layer = layer_index,
source = %shared_head_head_source,
vocab_size,
"qwen35 MTP loader: shared_head_head resolved"
);
let attn = load_mtp_attn(gguf, cfg, layer_index, device)?;
let (ffn, intermediate_size) = load_mtp_ffn(gguf, cfg, layer_index, device)?;
Ok(Some(MtpWeights {
layer_index,
hidden_size: cfg.hidden_size,
vocab_size,
intermediate_size,
loaded_tensor_names,
enorm,
hnorm,
eh_proj,
eh_proj_ggml_type,
embed_tokens,
shared_head_norm,
shared_head_head,
shared_head_head_ggml_type,
attn,
ffn,
}))
}
pub(crate) fn validate_mtp_tensor_topology(gguf: &GgufFile, cfg: &Qwen35Config) -> Result<()> {
if cfg.mtp_num_hidden_layers == 0 {
return Ok(());
}
ensure!(
cfg.mtp_num_hidden_layers == 1,
"qwen35 MTP loader supports exactly one nextn layer, got {}",
cfg.mtp_num_hidden_layers
);
let layer = cfg.num_hidden_layers;
ensure!(
!mtp_tensor_names(gguf, layer).is_empty(),
"qwen35 metadata advertises nextn_predict_layers=1 but no blk.{layer}.nextn.* or blk.{layer}.* MTP tensors were found"
);
let h = cfg.hidden_size as usize;
let q = (cfg.num_attention_heads as usize)
.checked_mul(cfg.head_dim as usize)
.context("MTP Q projection dimension overflow")?;
let kv = (cfg.num_key_value_heads as usize)
.checked_mul(cfg.head_dim as usize)
.context("MTP KV projection dimension overflow")?;
let d = cfg.head_dim as usize;
let p = format!("blk.{layer}");
let nextn = format!("{p}.nextn");
let require_shape = |name: &str, expected: &[usize]| -> Result<()> {
let info = gguf
.tensor_info(name)
.ok_or_else(|| anyhow::anyhow!("required MTP tensor `{name}` is missing"))?;
ensure!(
info.shape.as_slice() == expected,
"MTP tensor `{name}` shape {:?} != {expected:?}",
info.shape
);
Ok(())
};
for name in [
format!("{nextn}.enorm.weight"),
format!("{nextn}.hnorm.weight"),
format!("{nextn}.shared_head_norm.weight"),
format!("{p}.attn_norm.weight"),
format!("{p}.post_attention_norm.weight"),
] {
require_shape(&name, &[h])?;
}
let eh_name = format!("{nextn}.eh_proj.weight");
let eh_info = gguf
.tensor_info(&eh_name)
.ok_or_else(|| anyhow::anyhow!("required MTP tensor `{eh_name}` is missing"))?;
validate_native_projection_info(&eh_name, eh_info, h, 2 * h)?;
let embed = format!("{nextn}.embed_tokens.weight");
let embedding_rows = match (cfg.mtp_use_dedicated_embeddings, gguf.tensor_info(&embed)) {
(true, Some(info)) => {
ensure!(
info.shape.len() == 2 && info.shape[1] == h,
"MTP dedicated embedding tensor `{embed}` shape {:?} is not [vocab,{h}]",
info.shape
);
super::forward_gpu::validate_native_embedding_descriptor(
&embed,
info.ggml_type,
info.shape[0],
h,
info.byte_len,
)?;
info.shape[0]
}
(true, None) => bail!("MTP dedicated embedding tensor `{embed}` is missing"),
(false, Some(_)) => {
bail!("qwen35 MTP loader: mtp_use_dedicated_embeddings=False but `{embed}` is present")
}
(false, None) => {
gguf.tensor_info("token_embd.weight")
.context("MTP shared embedding tensor `token_embd.weight` is missing")?
.shape[0]
}
};
let head = format!("{nextn}.shared_head_head.weight");
let head_rows = if let Some(info) = gguf.tensor_info(&head) {
ensure!(
info.shape.len() == 2,
"MTP shared head tensor `{head}` is not a matrix"
);
validate_native_projection_info(&head, info, info.shape[0], h)?;
info.shape[0]
} else {
ensure!(
!cfg.mtp_use_dedicated_embeddings,
"qwen35 MTP loader: `{head}` is missing while dedicated embeddings are enabled"
);
let main_head = if gguf.tensor_info("output.weight").is_some() {
"output.weight"
} else {
"token_embd.weight"
};
let info = gguf
.tensor_info(main_head)
.ok_or_else(|| anyhow::anyhow!("MTP shared main head `{main_head}` is missing"))?;
ensure!(
info.shape.len() == 2,
"MTP shared main head `{main_head}` is not a matrix"
);
validate_native_projection_info(main_head, info, info.shape[0], h)?;
info.shape[0]
};
ensure!(
embedding_rows >= head_rows,
"MTP embedding rows {embedding_rows} cannot cover selected MTP head rows {head_rows}"
);
for (name, rows, cols) in [
(format!("{p}.attn_k.weight"), kv, h),
(format!("{p}.attn_v.weight"), kv, h),
(format!("{p}.attn_output.weight"), h, q),
] {
let info = gguf
.tensor_info(&name)
.ok_or_else(|| anyhow::anyhow!("required MTP tensor `{name}` is missing"))?;
validate_native_projection_info(&name, info, rows, cols)?;
}
require_shape(&format!("{p}.attn_q_norm.weight"), &[d])?;
require_shape(&format!("{p}.attn_k_norm.weight"), &[d])?;
let q_name = format!("{p}.attn_q.weight");
let q_info = gguf
.tensor_info(&q_name)
.ok_or_else(|| anyhow::anyhow!("required MTP tensor `{q_name}` is missing"))?;
ensure!(
q_info.shape.as_slice() == [q, h] || q_info.shape.as_slice() == [2 * q, h],
"MTP tensor `{q_name}` shape {:?} is neither [{q},{h}] nor [{},{h}]",
q_info.shape,
2 * q
);
validate_native_projection_info(&q_name, q_info, q_info.shape[0], h)?;
if q_info.shape.as_slice() == [q, h] {
let gate_name = format!("{p}.attn_gate.weight");
if let Some(gate_info) = gguf.tensor_info(&gate_name) {
validate_native_projection_info(&gate_name, gate_info, q, h)?;
}
} else {
ensure!(
gguf.tensor_info(&format!("{p}.attn_gate.weight")).is_none(),
"MTP fused Q+gate tensor cannot also have a split attn_gate tensor"
);
}
let has_dense = gguf.tensor_info(&format!("{p}.ffn_gate.weight")).is_some();
let has_moe = gguf
.tensor_info(&format!("{p}.ffn_gate_exps.weight"))
.is_some();
ensure!(
has_dense ^ has_moe,
"qwen35 MTP block {layer} must have exactly one dense or MoE FFN layout"
);
if has_dense {
let gate = gguf
.tensor_info(&format!("{p}.ffn_gate.weight"))
.expect("presence checked above");
ensure!(
gate.shape.len() == 2 && gate.shape[1] == h,
"MTP dense gate shape {:?} is not [intermediate,{h}]",
gate.shape
);
let intermediate = gate.shape[0];
require_shape(&format!("{p}.ffn_up.weight"), &[intermediate, h])?;
require_shape(&format!("{p}.ffn_down.weight"), &[h, intermediate])?;
let (gate_type, up_type, down_type) = dense_ffn_tensor_types(gguf, layer)?;
dense_ffn_storage(layer, gate_type, up_type, down_type)?;
} else {
let moe = cfg
.moe
.as_ref()
.context("MTP MoE tensors require MoE runtime metadata")?;
let experts = moe.num_experts as usize;
let intermediate = moe.moe_intermediate_size as usize;
let shared = moe.shared_expert_intermediate_size as usize;
require_shape(&format!("{p}.ffn_gate_inp.weight"), &[experts, h])?;
require_shape(
&format!("{p}.ffn_gate_exps.weight"),
&[experts, intermediate, h],
)?;
require_shape(
&format!("{p}.ffn_up_exps.weight"),
&[experts, intermediate, h],
)?;
require_shape(
&format!("{p}.ffn_down_exps.weight"),
&[experts, h, intermediate],
)?;
require_shape(&format!("{p}.ffn_gate_inp_shexp.weight"), &[1, h])?;
require_shape(&format!("{p}.ffn_gate_shexp.weight"), &[shared, h])?;
require_shape(&format!("{p}.ffn_up_shexp.weight"), &[shared, h])?;
require_shape(&format!("{p}.ffn_down_shexp.weight"), &[h, shared])?;
let gate_name = format!("{p}.ffn_gate_exps.weight");
let up_name = format!("{p}.ffn_up_exps.weight");
let down_name = format!("{p}.ffn_down_exps.weight");
let gate = gguf
.tensor_info(&gate_name)
.expect("MTP MoE gate presence checked above");
let up = gguf
.tensor_info(&up_name)
.expect("MTP MoE up shape checked above");
let down = gguf
.tensor_info(&down_name)
.expect("MTP MoE down shape checked above");
ensure!(
gate.ggml_type == up.ggml_type,
"MTP MoE gate/up quant types differ ({:?} vs {:?}); both buffers share one dispatch type",
gate.ggml_type,
up.ggml_type
);
ensure!(
qwen35_moe_expert_type_supported(gate.ggml_type),
"MTP MoE gate/up expert weights use unsupported {:?} storage",
gate.ggml_type
);
ensure!(
qwen35_moe_expert_type_supported(down.ggml_type),
"MTP MoE down expert weights use unsupported {:?} storage",
down.ggml_type
);
}
Ok(())
}
pub(super) fn mtp_tensor_names(gguf: &GgufFile, layer_index: u32) -> Vec<String> {
let p = format!("blk.{layer_index}.");
let nextn = format!("{p}nextn.");
let inner = [
"attn_norm.weight",
"post_attention_norm.weight",
"attn_q.weight",
"attn_gate.weight",
"attn_k.weight",
"attn_v.weight",
"attn_output.weight",
"attn_q_norm.weight",
"attn_k_norm.weight",
"ffn_gate.weight",
"ffn_up.weight",
"ffn_down.weight",
"ffn_gate_inp.weight",
"ffn_gate_exps.weight",
"ffn_up_exps.weight",
"ffn_down_exps.weight",
"ffn_gate_inp_shexp.weight",
"ffn_gate_shexp.weight",
"ffn_up_shexp.weight",
"ffn_down_shexp.weight",
];
let mut names = Vec::new();
for name in gguf.tensor_names() {
if name.starts_with(&nextn) || inner.iter().any(|suffix| name == format!("{p}{suffix}")) {
names.push(name.to_string());
}
}
names.sort();
names
}
fn load_mtp_attn(
gguf: &GgufFile,
cfg: &Qwen35Config,
layer_index: u32,
device: &MlxDevice,
) -> Result<MtpFullAttnWeightsGpu> {
let p = format!("blk.{layer_index}");
let h = cfg.hidden_size as usize;
let q_total = (cfg.num_attention_heads * cfg.head_dim) as usize;
let kv_total = (cfg.num_key_value_heads * cfg.head_dim) as usize;
let d = cfg.head_dim as usize;
let attn_norm = load_norm_gpu(gguf, &format!("{p}.attn_norm.weight"), h, device)?;
let post_attn_norm =
load_norm_gpu(gguf, &format!("{p}.post_attention_norm.weight"), h, device)?;
let q_name = format!("{p}.attn_q.weight");
let q_info = gguf
.tensor_info(&q_name)
.ok_or_else(|| anyhow::anyhow!("{q_name} not found"))?;
let q_gate = if q_info.shape.as_slice() == [q_total, h] {
let (q, q_type) = load_native_projection(gguf, &q_name, q_total, h, device)?;
let gate_name = format!("{p}.attn_gate.weight");
if gguf.tensor_info(&gate_name).is_some() {
let (gate, gate_type) = load_native_projection(gguf, &gate_name, q_total, h, device)?;
MtpQGateWeightsGpu::Gated(FullAttnQGateWeightsGpu::Split {
wq: q,
wq_ggml_type: q_type,
w_gate: gate,
w_gate_ggml_type: gate_type,
})
} else {
MtpQGateWeightsGpu::Ungated {
wq: q,
wq_ggml_type: q_type,
}
}
} else if q_info.shape.as_slice() == [2 * q_total, h] {
let (fused, q_type) = load_native_projection(gguf, &q_name, 2 * q_total, h, device)?;
MtpQGateWeightsGpu::Gated(FullAttnQGateWeightsGpu::Fused {
weight: fused,
ggml_type: q_type,
})
} else {
bail!(
"{q_name} shape {:?}, expected [{q_total},{h}] or [{}, {h}] interleaved Q+gate",
q_info.shape,
2 * q_total,
);
};
let (wk, wk_ggml_type) =
load_native_projection(gguf, &format!("{p}.attn_k.weight"), kv_total, h, device)?;
let (wv, wv_ggml_type) =
load_native_projection(gguf, &format!("{p}.attn_v.weight"), kv_total, h, device)?;
let (wo, wo_ggml_type) =
load_native_projection(gguf, &format!("{p}.attn_output.weight"), h, q_total, device)?;
Ok(MtpFullAttnWeightsGpu {
attn_norm,
post_attn_norm,
q_gate,
wk,
wk_ggml_type,
wv,
wv_ggml_type,
attn_q_norm: load_norm_gpu(gguf, &format!("{p}.attn_q_norm.weight"), d, device)?,
attn_k_norm: load_norm_gpu(gguf, &format!("{p}.attn_k_norm.weight"), d, device)?,
wo,
wo_ggml_type,
})
}
fn load_mtp_ffn(
gguf: &GgufFile,
cfg: &Qwen35Config,
layer_index: u32,
device: &MlxDevice,
) -> Result<(MtpFfnWeightsGpu, u32)> {
let p = format!("blk.{layer_index}");
let has_dense = gguf.tensor_info(&format!("{p}.ffn_gate.weight")).is_some();
let has_moe = gguf
.tensor_info(&format!("{p}.ffn_gate_exps.weight"))
.is_some();
match (has_dense, has_moe) {
(true, false) => load_mtp_dense_ffn(gguf, cfg, &p, device),
(false, true) => load_mtp_moe_ffn(gguf, cfg, layer_index, &p, device),
(true, true) => bail!(
"qwen35 MTP loader: block {layer_index} has BOTH dense (`{p}.ffn_gate.weight`) and \
MoE (`{p}.ffn_gate_exps.weight`) FFN tensors — GGUF is malformed"
),
(false, false) => bail!(
"qwen35 MTP loader: block {layer_index} has NEITHER dense (`{p}.ffn_gate.weight`) \
nor MoE (`{p}.ffn_gate_exps.weight`) inner FFN tensors"
),
}
}
fn load_mtp_dense_ffn(
gguf: &GgufFile,
cfg: &Qwen35Config,
p: &str,
device: &MlxDevice,
) -> Result<(MtpFfnWeightsGpu, u32)> {
let layer_idx = cfg.num_hidden_layers;
let (gate_type, up_type, down_type) = dense_ffn_tensor_types(gguf, layer_idx)?;
match dense_ffn_storage(layer_idx, gate_type, up_type, down_type)? {
DenseFfnStorage::Quantized => {
let weights_q = load_dense_ffn_quantized(gguf, layer_idx, cfg, device)
.with_context(|| format!("MTP native dense FFN {p}"))?;
let intermediate_size = weights_q.intermediate_size;
let dense_gpu = DenseFfnWeightsGpuQ::from_quantized(&weights_q);
Ok((
MtpFfnWeightsGpu::DenseQ { weights: dense_gpu },
intermediate_size,
))
}
DenseFfnStorage::Float => {
let h = cfg.hidden_size as usize;
let gate_name = format!("{p}.ffn_gate.weight");
let gate = load_f32_tensor(gguf, &gate_name, device)?;
ensure!(gate.len() % h == 0, "{gate_name} width mismatch");
let intermediate = gate.len() / h;
let up = load_f32_tensor(gguf, &format!("{p}.ffn_up.weight"), device)?;
let down = load_f32_tensor(gguf, &format!("{p}.ffn_down.weight"), device)?;
ensure!(
up.len() == intermediate * h,
"{p}.ffn_up.weight shape mismatch"
);
ensure!(
down.len() == h * intermediate,
"{p}.ffn_down.weight shape mismatch"
);
let weights =
DenseFfnWeightsGpu::from_cpu(&DenseFfnWeights { gate, up, down }, device)?;
let intermediate_size = intermediate as u32;
Ok((
MtpFfnWeightsGpu::Dense {
weights,
intermediate_size,
},
intermediate_size,
))
}
}
}
fn load_mtp_moe_ffn(
gguf: &GgufFile,
cfg: &Qwen35Config,
layer_index: u32,
p: &str,
device: &MlxDevice,
) -> Result<(MtpFfnWeightsGpu, u32)> {
let moe_cfg = cfg.moe.as_ref().ok_or_else(|| {
anyhow::anyhow!(
"qwen35 MTP loader: block {layer_index} has MoE FFN tensors (`{p}.ffn_gate_exps.weight`) \
but `cfg.moe` is None — the model metadata is inconsistent"
)
})?;
let weights_q = load_moe_ffn_quantized(gguf, layer_index, device)
.with_context(|| format!("MTP MoE FFN layer {layer_index}"))?;
let moe_gpu = MoeFfnWeightsGpuQ::from_quantized(
weights_q.expert_gate_q.clone(),
weights_q.expert_up_q.clone(),
weights_q.expert_down_q.clone(),
weights_q.ggml_type_gate_up,
weights_q.ggml_type_down,
moe_cfg.num_experts,
moe_cfg.moe_intermediate_size,
cfg.hidden_size,
&weights_q.router,
&weights_q.shared_gate_logit,
&weights_q.shared_gate,
&weights_q.shared_up,
&weights_q.shared_down,
device,
)
.with_context(|| format!("MTP MoE upload layer {layer_index}"))?;
let shape = MoeFfnShape {
hidden_size: cfg.hidden_size,
num_experts: moe_cfg.num_experts,
num_experts_per_tok: moe_cfg.num_experts_per_tok,
moe_intermediate_size: moe_cfg.moe_intermediate_size,
shared_intermediate_size: moe_cfg.shared_expert_intermediate_size,
};
Ok((
MtpFfnWeightsGpu::Moe {
weights: moe_gpu,
shape,
},
moe_cfg.moe_intermediate_size,
))
}
fn load_norm_gpu(gguf: &GgufFile, name: &str, len: usize, device: &MlxDevice) -> Result<MlxBuffer> {
let data = load_f32_tensor(gguf, name, device).with_context(|| name.to_string())?;
ensure!(data.len() == len, "{name} length {} != {len}", data.len());
upload_f32_weight(&data, device).with_context(|| format!("upload {name}"))
}