Struct MtpHead 

pub struct MtpHead { /* private fields */ }

2026-09-25: The MTP draft head.

Implementations

impl MtpHead

pub fn new( weights: MtpWeights, embed_tokens: DenseWeight, target_final_norm: DenseWeight, lm_head_nvfp4: QuantizedWeight, lm_head_nvfp4_t: Option<(QuantizedWeight, u32)>, config: &ModelConfig, gpu: &dyn GpuBackend, quant: MtpQuantization, mtp_vocab_size: u32, max_seq_len: usize, main_kv_blocks: usize, levers: &ModelLevers, ) -> Result<Self>

Trait Implementations

impl DraftProposer for MtpHead

fn circuit_draft( &self, config: &ModelConfig, levers: &ModelLevers, ) -> Option<DraftBinding>

2026-09-29: This proposer’s draft head as the circuit executor binds it; None for a proposer the circuit does not model.

fn set_circuit_draft(&self, runner: Option<Arc<dyn DraftRunner>>)

2026-09-29: Run runner in place of the draft head’s own forward from now on, or its own forward again with None. A proposer without a circuit binding ignores it.

fn alloc_state(&self, _gpu: &dyn GpuBackend) -> Result<Box<dyn ProposerState>>

2026-09-25: Allocate per-sequence proposer state.

fn propose( &self, last_token: u32, target_hidden: DevicePtr, position: usize, num_drafts: usize, state: &mut dyn ProposerState, ctx: &ForwardContext<'_>, stream: u64, draft_embed_target: Option<DevicePtr>, grammar_bitmask: Option<&[i32]>, _target_hidden_stack: Option<DevicePtr>, ) -> Result<Vec<u32>>

2026-09-25: Propose up to num_drafts tokens autoregressively.

fn propose_batch( &self, last_tokens: &[u32], target_hiddens: &[DevicePtr], positions: &[usize], num_drafts: usize, states: &mut [&mut dyn ProposerState], ctx: &ForwardContext<'_>, stream: u64, out_conf: Option<&mut Vec<Vec<f32>>>, ) -> Result<Option<Vec<Vec<u32>>>>

2026-09-25: Propose num_drafts tokens for each of last_tokens.len() sequences at once. Index i of every slice belongs to sequence i; target_hiddens[i] is its [1, hidden] BF16 hidden, and the pointers need not be contiguous. There is no grammar mask.

fn propose_batch_max( &self, buffers: &BufferArena, config: &ModelConfig, ) -> usize

2026-09-25: The widest batch Self::propose_batch accepts; 1 means per-sequence only. The scheduler groups batched proposes by it.

fn prefill_drafter( &self, prompt_tokens: &[u32], hiddens: DevicePtr, state: &mut dyn ProposerState, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<usize>

2026-09-25: Prefill the drafter’s KV over the prompt before the sequence’s first propose. hiddens is [P, hidden_size] BF16, row i the target’s hidden for prompt_tokens[i] before the final norm. Returns the drafter rows written; the default writes none.

fn drafter_rows(&self, state: &mut dyn ProposerState) -> usize

2026-09-25: Drafter KV rows held in state; 0 when unknown.

fn last_pair_key(&self, state: &mut dyn ProposerState) -> Option<usize>

2026-09-25: Sequence-space pair key of the newest drafter row; None when untracked. Rows are compacted, so the row count does not give the sequence position.

fn take_drafter_kv( &self, state: &mut dyn ProposerState, ) -> Option<(Vec<u32>, usize, Option<usize>)>

2026-09-25: Move this sequence’s drafter KV blocks out of its state, so free_state does not release them and the model can carry them to the next turn. Returns (blocks, rows, last_pair_key), or None when unsupported or empty. Afterwards the state must behave as freshly allocated.

fn install_drafter_kv( &self, state: &mut dyn ProposerState, blocks: Vec<u32>, rows: usize, last_pair_key: Option<usize>, ) -> bool

2026-09-25: Install carried blocks into a fresh state, the inverse of Self::take_drafter_kv. Returns false when unsupported; the caller then still owns the blocks and must free them.

fn free_drafter_kv(&self, blocks: &[u32])

2026-09-25: Release drafter KV blocks that no proposer state owns, such as a carried entry being replaced or dropped.

fn catchup_drafter( &self, tokens: &[u32], hiddens: DevicePtr, row_base: usize, pos_base: usize, state: &mut dyn ProposerState, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<usize>

2026-09-25: Append drafter rows at KV slots row_base .. with RoPE positions pos_base .. from (tokens, hiddens) pairs. Returns the rows written; the default writes none.

fn read_deferred_draft_token(&self, gpu: &dyn GpuBackend) -> Result<u32>

2026-09-25: The draft token id the last propose with draft_embed_target left on the device; the default returns 0.

fn catchup_batch( &self, tokens: &[Vec<u32>], hiddens: &[Vec<DevicePtr>], first_pos: &[usize], states: &mut [&mut dyn ProposerState], ctx: &ForwardContext<'_>, stream: u64, ) -> Result<usize>

2026-09-25: Append one drafter row per accepted draft of the last verify, for several sequences; the model calls it only with ModelLevers::mtp_kv_exact on. tokens[i] are sequence i’s accepted drafts, hiddens[i][k] the target hidden paired with tokens[i][k], and first_pos[i] the RoPE position of tokens[i][0]. Returns the rows written; the default writes none.

fn after_verify( &self, num_accepted: usize, state: &mut dyn ProposerState, _stream: u64, ) -> Result<()>

2026-09-25: Trim the proposer’s state after a verify that accepted num_accepted drafts.

fn free_state( &self, _gpu: &dyn GpuBackend, state: &mut dyn ProposerState, ) -> Result<()>

2026-09-25: Free per-sequence proposer state when a sequence finishes. DevicePtr has no Drop, so device memory alloc_state allocated leaks unless it is freed here; the default frees nothing.

fn alloc_state_for( &self, gpu: &dyn GpuBackend, budget_tokens: usize, ) -> Result<Box<dyn ProposerState>>

2026-09-25: Self::alloc_state given the sequence’s token budget (prompt length plus max_tokens), so state that scales with context can be sized to the request instead of --max-seq-len. usize::MAX means unknown. The default ignores the budget.

fn block_gamma(&self) -> Option<usize>

2026-09-25: The block size γ of a block-diffusion (DFlash) drafter; None for other proposers. With --dflash, serve sets the draft count to γ - 1, at least 1.

fn prefill_hidden_rows(&self, max_seq_len: usize) -> usize

2026-09-25: Rows of mtp_prefill_hidden this proposer can use, given the served --max-seq-len; the model sizes that [rows, hidden] BF16 buffer from it. The default is max_seq_len.

fn needs_comm(&self) -> bool

2026-09-25: Whether this proposer’s block is sharded across ranks, so its forward is given the communicator. The default is false; the GLM-5.3 MTP head returns mtp_ep_propose_enabled. When it is true on a multi-rank serve, the model sends the worker EP_CMD_MTP_PROPOSE before each propose, so both ranks run the same forward and collectives.

fn prefill_uses_shared_buffers(&self) -> bool

2026-09-25: Whether this proposer’s context prefill uses the shared forward scratch (ctx.buffers). If so, the model skips the prefill at the end of the target’s prefill, while the target still owns those buffers, and it runs at the first propose instead. Measured 2026-08-29 on GLM-5.3 (2x GB10): running it at the end of prefill changed the target’s output on 2 of 6 probes.