Struct MultiModuleMtpHead 

pub struct MultiModuleMtpHead { /* private fields */ }

2026-09-25: Independent MTP modules, one per draft slot.

Implementations

impl MultiModuleMtpHead

pub fn new(modules: Vec<MtpHead>) -> Result<Self>

2026-09-25: Assemble a multi-module proposer from per-module heads. Errors on an empty list.

pub fn num_modules(&self) -> usize

2026-09-25: Number of modules; propose caps num_drafts at it.

Trait Implementations

impl Debug for MultiModuleMtpHead

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl DraftProposer for MultiModuleMtpHead

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 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 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.

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 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 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 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 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.