Trait DraftProposer 

pub trait DraftProposer: Send + Sync {
    // Required methods
    fn alloc_state(
        &self,
        gpu: &dyn GpuBackend,
    ) -> Result<Box<dyn ProposerState>>;
    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>>;
    fn after_verify(
        &self,
        num_accepted: usize,
        state: &mut dyn ProposerState,
        stream: u64,
    ) -> Result<()>;

    // Provided methods
    fn alloc_state_for(
        &self,
        gpu: &dyn GpuBackend,
        budget_tokens: usize,
    ) -> Result<Box<dyn ProposerState>> { ... }
    fn block_gamma(&self) -> Option<usize> { ... }
    fn last_confidence(&self) -> Option<f32> { ... }
    fn prefill_hidden_rows(&self, max_seq_len: usize) -> usize { ... }
    fn needs_comm(&self) -> bool { ... }
    fn prefill_uses_shared_buffers(&self) -> bool { ... }
    fn drafter_rows(&self, _state: &mut dyn ProposerState) -> usize { ... }
    fn last_pair_key(&self, _state: &mut dyn ProposerState) -> Option<usize> { ... }
    fn take_drafter_kv(
        &self,
        _state: &mut dyn ProposerState,
    ) -> Option<(Vec<u32>, usize, Option<usize>)> { ... }
    fn install_drafter_kv(
        &self,
        _state: &mut dyn ProposerState,
        _blocks: Vec<u32>,
        _rows: usize,
        _last_pair_key: Option<usize>,
    ) -> bool { ... }
    fn free_drafter_kv(&self, _blocks: &[u32]) { ... }
    fn catchup_drafter(
        &self,
        _tokens: &[u32],
        _hiddens: DevicePtr,
        _row_base: usize,
        _pos_base: usize,
        _state: &mut dyn ProposerState,
        _ctx: &ForwardContext<'_>,
        _stream: u64,
    ) -> Result<usize> { ... }
    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>>>> { ... }
    fn propose_batch_max(
        &self,
        _buffers: &BufferArena,
        _config: &ModelConfig,
    ) -> usize { ... }
    fn prefill_drafter(
        &self,
        prompt_tokens: &[u32],
        hiddens: DevicePtr,
        state: &mut dyn ProposerState,
        ctx: &ForwardContext<'_>,
        stream: u64,
    ) -> Result<usize> { ... }
    fn read_deferred_draft_token(&self, gpu: &dyn GpuBackend) -> Result<u32> { ... }
    fn free_state(
        &self,
        gpu: &dyn GpuBackend,
        state: &mut dyn ProposerState,
    ) -> Result<()> { ... }
}

Required Methods

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.

Arguments

• last_token - The last verified token (target model output)
• target_hidden - Target model’s hidden states after final norm [1, hidden_size] BF16
• position - Current sequence position (for RoPE)
• num_drafts - Maximum number of draft tokens to produce
• state - Per-sequence proposer state
• ctx - Shared forward context (buffers, gpu, config)
• stream - CUDA stream handle
• grammar_bitmask - Optional XGrammar bitmask (ceil(vocab_size/32) i32 words). When Some, drafts are constrained to tokens the grammar accepts at the current matcher position; bit tok set ⇒ allowed. None preserves the unconstrained fast path.
• target_hidden_stack - Optional pointer to a contiguous buffer of 5 × target_hidden × bf16 containing the most-recently-decoded token’s hidden states captured at the drafter’s target_layer_ids (DFlash uses this; MTP ignores). Layout matches vLLM’s combine_hidden_states input: shallow-to-deep concatenation along the feature axis.

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

Called after target verification to trim proposer state.

num_accepted indicates how many draft tokens were accepted. The proposer should trim its KV cache / state to match.

Provided Methods

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.

usize::MAX means “unknown, use the ceiling”. Defaults to alloc_state, so proposers with fixed-size state need not implement 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 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 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.

The model allocates that buffer as [rows, hidden] BF16 before it knows anything about the proposer, so max_seq_len is the only bound it has — 4.0 GiB at 524,288 and h=4096. A proposer whose own architecture caps the position it can ever be asked for returns that cap instead, and the difference stops being allocated. See ANOMALIES A59 and the note in Glm5NextMtpHead::new.

🔴 Return a SMALLER number ONLY when the proposer can never be handed a position past it. A cap below the reachable context does not corrupt anything — the capture-coverage check at the propose site disables drafter-prefill for a sequence whose rows are short — but it silently costs acceptance on exactly the long prompts the feature exists for.

Default: max_seq_len, i.e. the pre-A59 sizing, which is correct for any proposer that can follow the target to the end of the served context.

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.

Default false, which is correct for the Qwen and DeepSeek-V4 drafters: their MTP modules load EVERY expert on EVERY rank, so the output is already complete and passing a comm would DOUBLE it via SUM.

🔴 GLM-5.3 is the opposite and it is not a choice: load_glm5next_mtp_module builds its MoE through the same Glm5NextMlpConfig the target layers use, so build_moe walks cfg.local_expert_range() and loads 144 of 288 experts; DsaTpPlan::new(tp_rank, tp_world_size, ..) splits the DSA heads the same way.

🪤 Returning true is NOT sufficient on its own — that is exactly what t58 did and it deadlocked at the first propose. The worker rank must ALSO execute the propose, or rank 0’s drafter collectives land against whatever the worker issues next. See EP_CMD_MTP_PROPOSE.

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.

MEASURED 2026-08-29 (GLM-5.3, 2x GB10, t61): the eager call site with the GLM drafter prefill engaged changed the TARGET’s completion on 2 of the 6 sealed probes and collapsed p1 from 0.625 to 0.045. The identical prefill work moved to the first propose is byte-identical on all six and takes p1 to 0.747. The call site is the only variable between the two arms; the exact colliding buffer is UNVERIFIED (norm_output and moe_output are the candidates the GLM block writes).

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

Row i of every slice belongs to sequence i; target_hiddens[i] is that sequence’s accepted-position hidden ([1, hidden] BF16, may be non-contiguous across i). Chains autoregressively per sequence like propose — position j uses (draft_{j-1}, drafter’s own hidden row i).

Returns Ok(None) when unsupported (caller falls back to the per-seq propose loop); Ok(Some(drafts)) with drafts[i].len() == num_drafts on success. Grammar-constrained sequences must not reach this path (callers gate on grammarless).

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.

Callers chunk by this instead of a hardcoded constant: a fixed cap of 4 made a 16-sequence step run 4 drafter forwards per position, each re-reading the whole drafter — the batched-propose lever’s own cost re-introduced by its caller.

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

• prompt_tokens — the prompt token ids t_0..t_{P-1}.
• hiddens — device buffer [P, hidden_size] BF16; row i is the target’s final-layer (pre-final-norm) hidden after processing t_i.

Returns the number of drafter positions written (0 = unsupported / already prefilled / nothing to do). Default: 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 free_state( &self, gpu: &dyn GpuBackend, state: &mut dyn ProposerState, ) -> Result<()>

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

Must be called when a sequence is finished to avoid resource leaks. gpu is threaded in (symmetric with alloc_state) so implementations can release raw device allocations stored on the state — DevicePtr has no Drop, so anything alloc_state allocated leaks unless it is explicitly freed here.

Implementors

impl DraftProposer for DeepseekV4MtpHead

impl DraftProposer for BlockDiffusionDraftHead

impl DraftProposer for Glm5NextMtpHead

impl DraftProposer for MtpHead

impl DraftProposer for MultiModuleMtpHead