metrale_model_arch/deepseek_v4_mtp/
proposer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `impl DraftProposer for DeepseekV4MtpHead`.
4//!
5//! Owner: model-arch, DeepSeek-V4 MTP.
6//! Invariants: none beyond the types.
7
8use super::*;
9
10impl DraftProposer for DeepseekV4MtpHead {
11    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn ProposerState>> {
12        Ok(Box::new(self.alloc_state_inner(gpu)?))
13    }
14
15    fn propose(
16        &self,
17        last_token: u32,
18        target_hidden: DevicePtr,
19        position: usize,
20        num_drafts: usize,
21        state: &mut dyn ProposerState,
22        ctx: &ForwardContext,
23        stream: u64,
24        _draft_embed_target: Option<DevicePtr>,
25        grammar_bitmask: Option<&[i32]>,
26        _target_hidden_stack: Option<DevicePtr>,
27    ) -> Result<Vec<u32>> {
28        let v4_state = state
29            .as_any_mut()
30            .downcast_mut::<DeepseekV4MtpProposerState>()
31            .ok_or_else(|| anyhow::anyhow!("Invalid V4 MTP proposer state"))?;
32
33        let mut drafts = Vec::with_capacity(num_drafts);
34        let mut current_token = last_token;
35        let mut current_hidden = target_hidden;
36        for i in 0..num_drafts {
37            if grammar_bitmask.is_some() && i > 0 {
38                tracing::warn!(
39                    "V4 MTP grammar-masked drafting with num_drafts>1 (i={i}); \
40                     mask held fixed across draft positions — acceptance may drop."
41                );
42            }
43            let draft = self.forward_one(
44                current_token,
45                current_hidden,
46                position + i,
47                v4_state,
48                ctx,
49                stream,
50                grammar_bitmask,
51            )?;
52            tracing::debug!(
53                "V4 MTP propose[{i}]: token={current_token} pos={} mtp_seq_len={} → draft={draft}",
54                position + i,
55                v4_state.seq_len,
56            );
57            drafts.push(draft);
58            current_token = draft;
59            // 2026-09-25: Later drafts feed on the MTP head's own collapsed
60            // hidden (`forward_one` leaves it in `hidden_states()`).
61            current_hidden = ctx.buffers.hidden_states();
62        }
63        v4_state.last_num_drafted = drafts.len();
64        Ok(drafts)
65    }
66
67    fn after_verify(
68        &self,
69        num_accepted: usize,
70        state: &mut dyn ProposerState,
71        _stream: u64,
72    ) -> Result<()> {
73        let v4_state = state
74            .as_any_mut()
75            .downcast_mut::<DeepseekV4MtpProposerState>()
76            .ok_or_else(|| anyhow::anyhow!("Invalid V4 MTP proposer state"))?;
77        // 2026-09-25: Trim `drafted - accepted` rejected entries (drafted
78        // counted as at least 1) from the MTP KV cache by rolling back
79        // `seq_len`; the next propose overwrites the slots.
80        let num_drafted = v4_state.last_num_drafted.max(1);
81        let num_to_trim = num_drafted.saturating_sub(num_accepted);
82        let old_sl = v4_state.seq_len;
83        if num_to_trim > 0 {
84            v4_state.seq_len = v4_state.seq_len.saturating_sub(num_to_trim);
85        }
86        tracing::debug!(
87            "V4 MTP after_verify: accepted={num_accepted} drafted={num_drafted} \
88             trim={num_to_trim} mtp_seq_len: {old_sl} → {}",
89            v4_state.seq_len,
90        );
91        Ok(())
92    }
93
94    fn free_state(&self, _gpu: &dyn GpuBackend, state: &mut dyn ProposerState) -> Result<()> {
95        let v4_state = state
96            .as_any_mut()
97            .downcast_mut::<DeepseekV4MtpProposerState>()
98            .ok_or_else(|| anyhow::anyhow!("Invalid V4 MTP proposer state"))?;
99        if !v4_state.block_table.is_empty() {
100            self.kv_cache.lock().free_blocks(&v4_state.block_table);
101            v4_state.block_table.clear();
102        }
103        v4_state.seq_len = 0;
104        Ok(())
105    }
106}
metrale_model_arch/deepseek_v4_mtp/proposer.rs

metrale_model_arch/deepseek_v4_mtp/
proposer.rs
