Struct MtpHead 

pub struct MtpHead { /* private fields */ }

MTP prediction head.

Implementations

impl MtpHead

pub fn new( weights: MtpWeights, embed_tokens: 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 alloc_state(&self, _gpu: &dyn GpuBackend) -> Result<Box<dyn ProposerState>>

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

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

Batched cross-sequence propose: draft num_drafts tokens for each of n = last_tokens.len() sequences, reading every drafter weight ONCE per draft position instead of once per sequence (the measured C=4 serialization: 12 x ~5 ms per-seq drafter forwards per batched verify step, ~62 ms of the ~180 ms step).

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

The widest batch Self::propose_batch can carry in ONE drafter forward per draft position, derived from this proposer’s resolved kernels and the arena’s row capacities. 1 = per-sequence only.

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

Prefill the drafter’s own context (KV cache) over the prompt, before the first propose() of a sequence (ATLAS_MTP_DRAFTER_PREFILL).

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

Current drafter KV length (rows), for the catch-up append point. 0 = unknown / not applicable (catch-up is skipped).

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

Sequence-space pair key of the newest drafter row (None = untracked; catch-up is skipped). The drafter row space is compacted, so rows cannot locate the drafter in the sequence — this can.

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

ATLAS_MTP_CARRY_DRAFTER: move this sequence’s drafter KV blocks OUT of its proposer state, so free_state releases nothing and the model can hold them for the next turn. Returns (blocks, rows, last_pair_key); None = unsupported or nothing to carry. After this call the state must behave as if freshly allocated.

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

Inverse of Self::take_drafter_kv: install carried blocks into a fresh proposer state. Returns false when unsupported (caller must then free the blocks itself).

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

Release drafter KV blocks that no proposer state owns (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>

Append drafter rows at KV slots row_base .. with RoPE positions pos_base .. from (tokens, hiddens) pairs — the catch-up feed. Returns rows written (0 = unsupported/no-op).

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

Read the draft token ID stored on GPU by the last propose() call that used draft_embed_target = Some(...). Returns 0 if not supported.

fn last_confidence(&self) -> Option<f32>

Chain confidence of the most recent propose (min top-1 softmax prob across its drafts), when the proposer computes it (draft_conf_tau > 0). None = not computed; callers must not gate on it then.

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

Called after target verification to trim proposer state.

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

Free per-sequence proposer state (KV cache blocks, device buffers, etc.).

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

Self::alloc_state with the sequence’s KNOWN token budget (prompt_len + max_tokens), so a proposer whose per-sequence state scales with context can size to what this request can actually reach instead of the global --max-seq-len ceiling. That distinction is what OOMs a high-concurrency long-context serve: the ceiling is per-sequence and paid n times, while a typical request needs a fraction of it.

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

The proposer’s trained block size γ, when it is a block-diffusion drafter (DFlash/DFlash2). The serve layer derives num_drafts from this — the head resolved it from the drafter checkpoint and is the SSOT. None = not a block drafter.

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

Rows this proposer can actually consume from mtp_prefill_hidden, given the served --max-seq-len.

fn needs_comm(&self) -> bool

True when this proposer’s block is SHARDED across ranks and its forward therefore needs the communicator (a routed-MoE all-reduce and a row-parallel o_proj reduce), like any target layer.

fn prefill_uses_shared_buffers(&self) -> bool

True when this proposer’s context prefill uses the SHARED forward scratch (ctx.buffers), so it must not run from the end-of-prefill eager hook — only from the first propose, where the target owns nothing.