Add the Multi-Token Prediction head and a model card

#1
by rkazants - opened
Optimum Intel Internal Testing org

Adds the Multi-Token Prediction (MTP) head of Qwen3.8-Flash-Next to the tiny model, as in the real checkpoint.

What changes

  • model.safetensors: 31 mtp.* tensors (one QSA decoder layer with hyper-connections and MoE, fc_embedding, fc_hidden, the pre-fc norms and the head's hyper_connection_mixer), named as in the real checkpoint.
  • config.json: the real mtp section (1 QSA layer) and mtp_num_hidden_layers 0 β†’ 1.
  • README.md: a model card describing the architecture, the initialization and the health checks.

Compatibility

  • The 243 main-model tensors are bit-identical to the current version, so existing tests are unaffected.
  • transformers ignores mtp.* when loading (_keys_to_ignore_on_load_unexpected).

Initialization and health checks

  • The MTP weights use the same init scheme as the main model: layers std 0.25, router std 1.
  • Checked on text, image and video prompts with a reference MTP forward that follows vLLM qwen4_exp:
    • logits are finite, with std β‰ˆ 4 and median greedy top-2 margins of 0.55–0.82;
    • 73–100% of positions get distinct predictions;
    • all experts are used, and its QSA drops 23–93% of the visible tokens;
    • the head depends on both of its inputs (next-token embedding and hidden streams);
    • a chained second draft step stays healthy.

πŸ€– Generated with Claude Code

rkazants changed pull request status to merged

Sign up or log in to comment