Qwen35MtpModel

Model type

Qwen35MtpModel

Class

Qwen35MtpModel

Task

qwen35-mtp

Source

models/qwen35_mtp.py

Description

Qwen3.6 MTP self-speculative head (a single cross-conditioned full-attention block).

Inputs (graph-level, set up by :class:~mobius.tasks.Qwen35MtpTask): input_ids or inputs_embeds: Dedicated embeddings consume token IDs; fallback embeddings arrive from the target as [batch, seq_len, hidden]. hidden_states: [batch, seq_len, hidden] (model dtype) — the target model’s last hidden state h_i (post-final-norm). attention_mask: [batch, total_seq_len] INT64. position_ids: [batch, seq_len] INT64. past_key_values: standard GQA KV cache for the single MTP layer.

Outputs: logits or mtp_hidden: Dedicated heads emit vocabulary logits; fallback heads emit final hidden states for the target’s shared lm_head. present_key_values: updated KV cache for the single MTP layer.

Usage

mobius build --model <MODEL_ID> output_dir/
from mobius import build

model = build("<MODEL_ID>")