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

self.max_seq_len is ALREADY capped at max_dsa_context by new, so this both rightsizes the model’s capture buffer and keeps it in lockstep with the drafter’s own bounds checks — a capture longer than the drafter’s row space could never be read. ANOMALIES A59.

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

Release everything alloc_state allocated.

Without this the head inherits DraftProposer::free_state’s no-op default, whose own doc says: “DevicePtr has no Drop, so anything alloc_state allocated leaks unless it is explicitly freed here.” That is exactly what happened — every finished sequence leaked its indexer cache. The cache is sized from serve_max_seq_len, so the leak scales with --max-seq-len: ~806 MB per sequence at --max-seq-len 131072, which walks a unified-memory host into the ground in a handful of requests (ANOMALIES A75). DeepseekV4MtpHead and MultiModuleMtp already override this; the GLM port did not.

🔴 Invariant L2 (slot reuse), not a line order: when this slot is re-occupied its decode_graph and verify2/3/4_graph — which bake these exact pointers — must already be destroyed AND these pointers freed and nulled. free_sequence satisfies both. ANOMALIES A56 is the history; the invariant is slot reuse, not the order of the two blocks. (The released flag below is what makes a second call safe.)

fn needs_comm(&self) -> bool

🔴 EP-sharded MoE (144 of 288 experts) + row-parallel DSA o_proj. Without the communicator this block drafts from HALF of both. See the trait doc for why that is only safe once the WORKER rank runs propose too.

fn prefill_uses_shared_buffers(&self) -> bool

🔴 The GLM context prefill runs the block through ctx.buffers. See the trait doc — running it from the end-of-prefill hook corrupts the TARGET’s output.

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

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>

🪤 pos_base is ignored: this drafter’s RoPE position is its slot (see rows_impl), so the caller’s sequence-space position is already row_base up to the uniform shift. A feed that does not start exactly at drafter_rows() is refused by rows_impl.

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

Allocate per-sequence proposer state.

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 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 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 after_verify( &self, num_accepted: usize, state: &mut dyn ProposerState, _stream: u64, ) -> Result<()>

Called after target verification to trim proposer state.

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