Module mtp_head 

2026-09-25: The MTP (multi-token prediction) draft head, a DraftProposer. One decoder layer drafts from the token embedding and the target’s hidden: pre-fc norms, concat, fc, gated attention over its own KV cache, MoE or dense FFN, final norm, the NVFP4 LM head and argmax. MtpQuantization selects the projection weights’ precision.

Owner: model-layers (MTP head). Invariants:

• mtp_rows_to_trim never returns more than num_drafted.

Structs

MtpHead

2026-09-25: The MTP draft head.

MtpProposerState

2026-09-25: Per-sequence MTP proposer state.

Enums

MtpQuantization

2026-09-25: Precision of the MTP head’s projection weights.

Functions

mtp_drafter_prefill_enabled

2026-09-25: Whether the drafter prefill is on (levers.drafter.prefill, on unless METRALE_NO_MTP_DRAFTER_CONTEXT turns it off; policy in crate::drafter_context). The prefill builds the drafter’s KV over the whole prompt before the first propose. The drafter’s K/V depend only on its input pair, not on attention, so it runs the norms, fc, K/V projections, RoPE and the cache write, with no attention pass.

mtp_experts_nvfp4

2026-10-02: Whether a BF16 MoE MTP head drafts on NVFP4 copies of its experts. Off unless the serve published otherwise: the head’s declared BF16 is the default.

set_mtp_experts_nvfp4_from_cli

2026-10-02: Publish --mtp-experts-nvfp4; returns the value in force, which differs when something read it first.