Module mtp_head 

MTP (Multi-Token Prediction) head implementing DraftProposer.

Single transformer decoder layer trained jointly with the target model. Forward pass: embed+hidden concat → fc → attention → MoE → norm → lm_head → argmax.

Weight precision is parameterized via MtpQuantization: NVFP4 (4-bit), FP8 (8-bit), or BF16 (16-bit). Higher precision improves draft acceptance at the cost of increased MTP forward latency.

Structs

MtpHead

MTP prediction head.

MtpProposerState

Per-sequence MTP proposer state.

Enums

MtpQuantization

MTP head weight precision.

Functions

mtp_drafter_prefill_enabled

Drafter context prefill — ON by default, cached once.