metrale_model_layers/layers/mtp_head/
draft_proposer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `DraftProposer` for [`MtpHead`]: single and batched propose, drafter-KV
4//! hand-over (take, install, free), catch-up rows and the trim after verify.
5//!
6//! Owner: model-layers (MTP head).
7//! Invariants:
8//! - `install_drafter_kv` installs only into a state with no blocks and no rows.
9//! - `after_verify` trims at most `max(last_num_drafted, 1)` rows.
10use super::*;
11
12impl DraftProposer for MtpHead {
13    fn circuit_draft(
14        &self,
15        config: &metrale_config::ModelConfig,
16        levers: &crate::layers::ops::ModelLevers,
17    ) -> Option<crate::circuit_exec::DraftBinding> {
18        Some(self.circuit_binding(config, levers))
19    }
20
21    fn set_circuit_draft(
22        &self,
23        runner: Option<std::sync::Arc<dyn crate::circuit_exec::DraftRunner>>,
24    ) {
25        self.install_circuit(runner);
26    }
27
28    fn alloc_state(&self, _gpu: &dyn GpuBackend) -> Result<Box<dyn ProposerState>> {
29        Ok(Box::new(MtpProposerState {
30            block_table: Vec::new(),
31            seq_len: 0,
32            last_num_drafted: 0,
33            last_pair_key: None,
34            last_drafts: Vec::new(),
35            pending_catchup: Vec::new(),
36        }))
37    }
38
39    fn propose(
40        &self,
41        last_token: u32,
42        target_hidden: DevicePtr,
43        position: usize,
44        num_drafts: usize,
45        state: &mut dyn ProposerState,
46        ctx: &ForwardContext,
47        stream: u64,
48        draft_embed_target: Option<DevicePtr>,
49        grammar_bitmask: Option<&[i32]>,
50        _target_hidden_stack: Option<DevicePtr>,
51    ) -> Result<Vec<u32>> {
52        let mtp_state = state
53            .as_any_mut()
54            .downcast_mut::<MtpProposerState>()
55            .ok_or_else(|| anyhow::anyhow!("Invalid MTP proposer state"))?;
56
57        // 2026-10-02: `--draft-confidence-stop`: the chain stops after the first
58        // draft whose top-1 probability is below tau (`speculative::draft_stop`).
59        let stop_lp = crate::speculative::draft_stop::draft_stop_logprob();
60        let mut drafts = Vec::with_capacity(num_drafts);
61        let mut current_token = last_token;
62        let mut current_hidden = target_hidden;
63
64        for i in 0..num_drafts {
65            let embed_target = if i == num_drafts - 1 {
66                draft_embed_target
67            } else {
68                None
69            };
70            // 2026-09-25: The grammar matcher is not advanced between draft positions:
71            // every position gets the same mask, with a warning past the first.
72            if grammar_bitmask.is_some() && i > 0 {
73                tracing::warn!(
74                    "MTP grammar-masked drafting called with num_drafts>1 (i={i}); \
75                     mask held fixed across draft positions — acceptance may drop."
76                );
77            }
78            let mask_for_draft = grammar_bitmask;
79            let draft = self.forward_one(
80                current_token,
81                current_hidden,
82                position + i,
83                mtp_state,
84                ctx,
85                stream,
86                embed_target,
87                mask_for_draft,
88                i == 0,
89            )?;
90            tracing::debug!(
91                "MTP propose[{i}]: token={current_token} pos={} mtp_seq_len={} → draft={draft}",
92                position + i,
93                mtp_state.seq_len,
94            );
95            drafts.push(draft);
96            current_token = draft;
97            // 2026-09-25: Later drafts read the drafter's own hidden (`chain_hidden`).
98            current_hidden = Self::chain_hidden(ctx);
99            // 2026-10-02: The last draft's log-probability is NaN when it was not
100            // measured (a masked or deferred draft), which never stops the chain.
101            if let Some(ln_tau) = stop_lp {
102                let lp = f32::from_bits(
103                    self.last_conf_bits
104                        .load(std::sync::atomic::Ordering::Relaxed),
105                );
106                if !lp.is_nan() && !crate::speculative::draft_stop::chain_continues(lp, ln_tau) {
107                    break;
108                }
109            }
110        }
111
112        mtp_state.last_num_drafted = drafts.len();
113        mtp_state.last_drafts.clone_from(&drafts);
114        Ok(drafts)
115    }
116
117    fn propose_batch(
118        &self,
119        last_tokens: &[u32],
120        target_hiddens: &[DevicePtr],
121        positions: &[usize],
122        num_drafts: usize,
123        states: &mut [&mut dyn ProposerState],
124        ctx: &ForwardContext,
125        stream: u64,
126        out_conf: Option<&mut Vec<Vec<f32>>>,
127    ) -> Result<Option<Vec<Vec<u32>>>> {
128        if !self.can_propose_batch(last_tokens.len(), ctx.buffers, ctx.config) {
129            return Ok(None);
130        }
131        // 2026-09-30: Under `--forward circuit` the batch runs the circuit's n-row draft
132        // program, which writes the D-Cut confidences; a width it has no program for, or a
133        // caller without D-Cut (whose legacy argmax the plan does not state), drafts each
134        // sequence alone on the single-row program (`forward.rs`).
135        let circuit = self.circuit_draft.read().clone();
136        if let Some(r) = &circuit
137            && !(r.serves(last_tokens.len() as u64) && out_conf.is_some())
138        {
139            return Ok(None);
140        }
141        // 2026-09-25: The width policy sees only the config and the arena, so the
142        // grouped FP8 decode's context terms (kill switch, FP32 routing, EP) are
143        // checked here: a refusal falls back per sequence instead of failing
144        // mid-chain, and is logged once per process.
145        if let Some(moe) = self.moe_grouped.as_ref()
146            && !moe.any_grouped_decode_ok(last_tokens.len(), ctx)
147        {
148            static LOGGED: std::sync::Once = std::sync::Once::new();
149            LOGGED.call_once(|| {
150                tracing::warn!(
151                    "MTP propose_batch: grouped FP8 MoE decode refused for n={} (kill switch \
152                     METRALE_NO_FP8_MOE_GROUPED_DECODE, FP32 routing or EP) — the MoE drafter \
153                     falls back to per-sequence propose",
154                    last_tokens.len()
155                );
156            });
157            return Ok(None);
158        }
159        // 2026-09-25: A state that is not an `MtpProposerState` makes the whole batch
160        // fall back per sequence.
161        let mut mtp_states: Vec<&mut MtpProposerState> = Vec::with_capacity(states.len());
162        for s in states.iter_mut() {
163            match s.as_any_mut().downcast_mut::<MtpProposerState>() {
164                Some(st) => mtp_states.push(st),
165                None => return Ok(None),
166            }
167        }
168        self.propose_batch_impl(
169            circuit.as_deref(),
170            last_tokens,
171            target_hiddens,
172            positions,
173            num_drafts,
174            &mut mtp_states,
175            ctx,
176            stream,
177            out_conf,
178        )
179        .map(Some)
180    }
181
182    fn propose_batch_max(
183        &self,
184        buffers: &metrale_gpu_runtime::buffers::BufferArena,
185        config: &metrale_config::ModelConfig,
186    ) -> usize {
187        MtpHead::propose_batch_max(self, buffers, config)
188    }
189
190    fn prefill_drafter(
191        &self,
192        prompt_tokens: &[u32],
193        hiddens: DevicePtr,
194        state: &mut dyn ProposerState,
195        ctx: &ForwardContext,
196        stream: u64,
197    ) -> Result<usize> {
198        self.prefill_drafter_impl(prompt_tokens, hiddens, state, ctx, stream)
199    }
200
201    fn drafter_rows(&self, state: &mut dyn ProposerState) -> usize {
202        state
203            .as_any_mut()
204            .downcast_mut::<MtpProposerState>()
205            .map(|s| s.seq_len)
206            .unwrap_or(0)
207    }
208
209    fn last_pair_key(&self, state: &mut dyn ProposerState) -> Option<usize> {
210        state
211            .as_any_mut()
212            .downcast_mut::<MtpProposerState>()
213            .and_then(|s| s.last_pair_key)
214    }
215
216    fn take_drafter_kv(
217        &self,
218        state: &mut dyn ProposerState,
219    ) -> Option<(Vec<u32>, usize, Option<usize>)> {
220        let st = state.as_any_mut().downcast_mut::<MtpProposerState>()?;
221        if st.block_table.is_empty() || st.seq_len == 0 {
222            return None;
223        }
224        let blocks = std::mem::take(&mut st.block_table);
225        let rows = st.seq_len;
226        let key = st.last_pair_key;
227        // 2026-09-25: No blocks, no rows, no pair key: `free_state` then has nothing to
228        // release.
229        st.seq_len = 0;
230        st.last_pair_key = None;
231        st.last_num_drafted = 0;
232        Some((blocks, rows, key))
233    }
234
235    fn install_drafter_kv(
236        &self,
237        state: &mut dyn ProposerState,
238        blocks: Vec<u32>,
239        rows: usize,
240        last_pair_key: Option<usize>,
241    ) -> bool {
242        let Some(st) = state.as_any_mut().downcast_mut::<MtpProposerState>() else {
243            return false;
244        };
245        // 2026-09-25: Only into an empty state; otherwise its blocks would leak.
246        if !st.block_table.is_empty() || st.seq_len != 0 {
247            return false;
248        }
249        st.block_table = blocks;
250        st.seq_len = rows;
251        st.last_pair_key = last_pair_key;
252        true
253    }
254
255    fn free_drafter_kv(&self, blocks: &[u32]) {
256        if !blocks.is_empty() {
257            self.kv_cache.lock().free_blocks(blocks);
258        }
259    }
260
261    fn catchup_drafter(
262        &self,
263        tokens: &[u32],
264        hiddens: DevicePtr,
265        row_base: usize,
266        pos_base: usize,
267        state: &mut dyn ProposerState,
268        ctx: &ForwardContext,
269        stream: u64,
270    ) -> Result<usize> {
271        self.drafter_rows_impl(tokens, hiddens, row_base, pos_base, state, ctx, stream)
272    }
273
274    fn read_deferred_draft_token(&self, gpu: &dyn GpuBackend) -> Result<u32> {
275        self.read_deferred_draft_token(gpu)
276    }
277
278    fn catchup_batch(
279        &self,
280        tokens: &[Vec<u32>],
281        hiddens: &[Vec<DevicePtr>],
282        first_pos: &[usize],
283        states: &mut [&mut dyn ProposerState],
284        ctx: &ForwardContext,
285        stream: u64,
286    ) -> Result<usize> {
287        self.catchup_batch_impl(tokens, hiddens, first_pos, states, ctx, stream)
288    }
289
290    fn after_verify(
291        &self,
292        num_accepted: usize,
293        state: &mut dyn ProposerState,
294        _stream: u64,
295    ) -> Result<()> {
296        let mtp_state = state
297            .as_any_mut()
298            .downcast_mut::<MtpProposerState>()
299            .ok_or_else(|| anyhow::anyhow!("Invalid MTP proposer state"))?;
300
301        let num_drafted = mtp_state.last_num_drafted.max(1);
302        // 2026-09-25: With exact KV, row 0 was built from the target hidden and stays;
303        // rows 1.. were built from drafter hiddens and go, and the catch-up
304        // rebuilds the accepted ones. Otherwise `mtp_rows_to_trim` decides.
305        let num_to_trim = if self.kv_exact {
306            num_drafted - 1
307        } else {
308            mtp_rows_to_trim(
309                num_drafted,
310                num_accepted,
311                crate::speculative::mtp_refeed_accepted_enabled(),
312            )
313        };
314        let old_sl = mtp_state.seq_len;
315        if num_to_trim > 0 {
316            mtp_state.seq_len = mtp_state.seq_len.saturating_sub(num_to_trim);
317            // 2026-09-25: Trimmed rows have consecutive pair keys, so the newest
318            // surviving key moves back by the same count.
319            if let Some(k) = mtp_state.last_pair_key {
320                mtp_state.last_pair_key = Some(k.saturating_sub(num_to_trim));
321            }
322        }
323        tracing::debug!(
324            "MTP after_verify: accepted={num_accepted} drafted={num_drafted} trim={num_to_trim} mtp_seq_len: {old_sl} → {}",
325            mtp_state.seq_len,
326        );
327        Ok(())
328    }
329
330    fn free_state(&self, _gpu: &dyn GpuBackend, state: &mut dyn ProposerState) -> Result<()> {
331        let mtp_state = state
332            .as_any_mut()
333            .downcast_mut::<MtpProposerState>()
334            .ok_or_else(|| anyhow::anyhow!("Invalid MTP proposer state"))?;
335        if !mtp_state.block_table.is_empty() {
336            self.kv_cache.lock().free_blocks(&mtp_state.block_table);
337            mtp_state.block_table.clear();
338        }
339        mtp_state.seq_len = 0;
340        Ok(())
341    }
342}
metrale_model_layers/layers/mtp_head/draft_proposer.rs

metrale_model_layers/layers/mtp_head/
draft_proposer.rs
