Struct Glm5NextMtpHead 

pub struct Glm5NextMtpHead { /* private fields */ }

Implementations

impl Glm5NextMtpHead

pub fn new( module: Glm5NextMtpModule, embed_tokens: DenseWeight, lm_head: DenseWeight, config: &ModelConfig, gpu: &dyn GpuBackend, max_seq_len: usize, ) -> Result<Self>

Trait Implementations

impl DraftProposer for Glm5NextMtpHead

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

2026-09-25: self.max_seq_len is already capped at max_dsa_context by new, so the model’s capture buffer is no longer than the drafter’s own row space.

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

2026-09-25: Release everything alloc_state allocated: the DSA indexer cache and the scratch buffers. DevicePtr has no Drop, so without this override they would leak. A second call is a no-op (released).

fn needs_comm(&self) -> bool

2026-09-25: The block’s routed experts are split across EP ranks and its DSA o_proj is row-parallel, so its forward needs the communicator. On unless METRALE_NO_MTP_EP_PROPOSE=1 (mtp_ep_propose_enabled).

fn prefill_uses_shared_buffers(&self) -> bool

2026-09-25: True: a context row normalises into ctx.buffers (drafter_write_kv_row).

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

2026-09-25: Dense row space: the newest row’s slot is its pair key.

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: pos_base is unused: this drafter’s RoPE position is its KV slot (see rows_impl). A feed that does not start exactly at drafter_rows() writes nothing.

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

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

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

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

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 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 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 alloc_state_for( &self, gpu: &dyn GpuBackend, budget_tokens: usize, ) -> Result<Box<dyn ProposerState>, Error>

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 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 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>>>, Error>

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, Error>

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 read_deferred_draft_token(&self, gpu: &dyn GpuBackend) -> Result<u32, Error>

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

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.