metrale_model_arch/glm5next_mtp_head/
proposer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `impl DraftProposer for Glm5NextMtpHead`: state lifetime, context prefill, propose and rollback.
4//!
5//! Owner: model-arch (GLM-5.3 MTP drafter).
6//! Invariants:
7//! - `after_verify` never trims the first row of a propose.
8//! - A second `free_state` on the same state is a no-op.
9
10use super::*;
11
12impl DraftProposer for Glm5NextMtpHead {
13    /// 2026-09-25: `self.max_seq_len` is already capped at `max_dsa_context` by `new`, so the
14    /// model's capture buffer is no longer than the drafter's own row space.
15    fn prefill_hidden_rows(&self, max_seq_len: usize) -> usize {
16        max_seq_len.min(self.max_seq_len)
17    }
18
19    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn ProposerState>> {
20        let dsa = match &self.module.layer.mixer {
21            crate::glm5next_layer::Glm5NextMixer::Dsa(l) => Glm5NextDsaState::alloc(gpu, &l.cfg)?,
22            _ => bail!("GLM MTP block is not a DSA layer"),
23        };
24        let h = self.hidden;
25        // 2026-09-25: Claims every block of the private pool, which `new` sizes for one
26        // sequence of `max_seq_len` rows.
27        let blocks = (self.max_seq_len / 16 + 2) as u32;
28        Ok(Box::new(Glm5NextMtpProposerState {
29            dsa,
30            seq_len: 0,
31            block_table: (0..blocks).collect(),
32            last_drafted: 0,
33            concat: gpu.alloc(2 * h * 2)?,
34            x: gpu.alloc(h * 2)?,
35            logits: gpu.alloc(self.vocab * 2)?,
36            arg: gpu.alloc(4)?,
37            head_xchg: gpu.alloc(16)?,
38            released: false,
39        }))
40    }
41
42    /// 2026-09-25: Release everything `alloc_state` allocated: the DSA indexer cache and the
43    /// scratch buffers. `DevicePtr` has no `Drop`, so without this override they would leak.
44    /// A second call is a no-op (`released`).
45    fn free_state(&self, gpu: &dyn GpuBackend, state: &mut dyn ProposerState) -> Result<()> {
46        let st = state
47            .as_any_mut()
48            .downcast_mut::<Glm5NextMtpProposerState>()
49            .ok_or_else(|| anyhow::anyhow!("Invalid GLM MTP proposer state"))?;
50        if st.released {
51            return Ok(());
52        }
53        st.released = true;
54        st.dsa.free(gpu)?;
55        for p in [st.concat, st.x, st.logits, st.arg, st.head_xchg] {
56            gpu.free(p)?;
57        }
58        // 2026-09-25: The drafter's private pool is claimed whole by `alloc_state`
59        // (`(0..blocks).collect()`), not drawn from an allocator, so there is
60        // nothing to hand back; clearing it stops a freed state from looking live.
61        st.block_table.clear();
62        st.seq_len = 0;
63        Ok(())
64    }
65
66    /// 2026-09-25: The block's routed experts are split across EP ranks and its DSA `o_proj`
67    /// is row-parallel, so its forward needs the communicator. On unless
68    /// `METRALE_NO_MTP_EP_PROPOSE=1` (`mtp_ep_propose_enabled`).
69    fn needs_comm(&self) -> bool {
70        metrale_model_layers::speculative::mtp_ep_propose_enabled()
71    }
72
73    /// 2026-09-25: True: a context row normalises into `ctx.buffers` (`drafter_write_kv_row`).
74    fn prefill_uses_shared_buffers(&self) -> bool {
75        true
76    }
77
78    fn drafter_rows(&self, state: &mut dyn ProposerState) -> usize {
79        state
80            .as_any_mut()
81            .downcast_mut::<Glm5NextMtpProposerState>()
82            .map_or(0, |st| st.seq_len)
83    }
84
85    /// 2026-09-25: Dense row space: the newest row's slot is its pair key.
86    fn last_pair_key(&self, state: &mut dyn ProposerState) -> Option<usize> {
87        state
88            .as_any_mut()
89            .downcast_mut::<Glm5NextMtpProposerState>()
90            .and_then(|st| st.seq_len.checked_sub(1))
91    }
92
93    fn prefill_drafter(
94        &self,
95        prompt_tokens: &[u32],
96        hiddens: DevicePtr,
97        state: &mut dyn ProposerState,
98        ctx: &ForwardContext,
99        stream: u64,
100    ) -> Result<usize> {
101        let t0 = std::time::Instant::now();
102        let rows = self.rows_impl(prompt_tokens, hiddens, 0, state, ctx, stream)?;
103        // 2026-09-25: `rows_impl` returns 0 without work once the drafter has rows, so only a
104        // call that wrote rows logs.
105        if rows > 0 {
106            tracing::info!(
107                "GLM MTP drafter prefill: {rows} rows ({} prompt tokens) in {:.1} ms",
108                prompt_tokens.len(),
109                t0.elapsed().as_secs_f64() * 1e3,
110            );
111        }
112        Ok(rows)
113    }
114
115    /// 2026-09-25: `pos_base` is unused: this drafter's RoPE position is its KV slot (see
116    /// `rows_impl`). A feed that does not start exactly at `drafter_rows()` writes nothing.
117    fn catchup_drafter(
118        &self,
119        tokens: &[u32],
120        hiddens: DevicePtr,
121        row_base: usize,
122        _pos_base: usize,
123        state: &mut dyn ProposerState,
124        ctx: &ForwardContext,
125        stream: u64,
126    ) -> Result<usize> {
127        self.rows_impl(tokens, hiddens, row_base, state, ctx, stream)
128    }
129
130    #[allow(clippy::too_many_arguments)]
131    fn propose(
132        &self,
133        last_token: u32,
134        target_hidden: DevicePtr,
135        position: usize,
136        num_drafts: usize,
137        state: &mut dyn ProposerState,
138        ctx: &ForwardContext,
139        stream: u64,
140        _draft_embed_target: Option<DevicePtr>,
141        _grammar_bitmask: Option<&[i32]>,
142        _target_hidden_stack: Option<DevicePtr>,
143    ) -> Result<Vec<u32>> {
144        let st = state
145            .as_any_mut()
146            .downcast_mut::<Glm5NextMtpProposerState>()
147            .ok_or_else(|| anyhow::anyhow!("not a GLM MTP proposer state"))?;
148        // 2026-09-25: A drafter ahead of the target's `position` is rewound to it, so its
149        // indexer does not select over rows past the target's context.
150        if st.seq_len > position {
151            st.dsa.rewind_to(position)?;
152            st.seq_len = position;
153        }
154        if metrale_model_layers::speculative::mtp_refeed_debug() {
155            let fp = metrale_model_layers::speculative::hidden_fingerprint(
156                ctx.gpu,
157                target_hidden,
158                self.hidden,
159            );
160            tracing::info!(
161                "GLM_MTP_DBG propose position={position} drafter_rows={} tok={last_token} \
162                 fp_target={fp:016x}",
163                st.seq_len,
164            );
165        }
166        let mut drafts = Vec::with_capacity(num_drafts);
167        let mut token = last_token;
168        let mut hidden = target_hidden;
169        for i in 0..num_drafts {
170            let d = self.forward_one(token, hidden, position + i, st, ctx, stream)?;
171            drafts.push(d);
172            token = d;
173            // 2026-09-25: Draft 1 consumes the target's verified hidden; every later draft
174            // consumes the drafter's own block output.
175            hidden = st.x;
176        }
177        st.last_drafted = drafts.len();
178        Ok(drafts)
179    }
180
181    fn after_verify(
182        &self,
183        num_accepted: usize,
184        state: &mut dyn ProposerState,
185        _stream: u64,
186    ) -> Result<()> {
187        let st = state
188            .as_any_mut()
189            .downcast_mut::<Glm5NextMtpProposerState>()
190            .ok_or_else(|| anyhow::anyhow!("not a GLM MTP proposer state"))?;
191        // 2026-09-25: Rejected rows are unreachable: the indexer reads `[0, len)` and the next
192        // propose writes from `seq_len`, so rolling the counters back is the whole rollback.
193        //
194        // The first row of a propose is always kept. It pairs the last committed token with the
195        // target's own hidden, both known at propose time, so a rejected draft does not make it
196        // wrong. Only later rows depend on a draft having been accepted, and trimming the first
197        // would drop a real row from the dense row space and shift every later RoPE position.
198        // At `num_drafts = 1` nothing is trimmed.
199        let keep = st.last_drafted.min(num_accepted + 1);
200        let trim = st.last_drafted - keep;
201        if trim > 0 {
202            st.seq_len = st.seq_len.saturating_sub(trim);
203            st.dsa.rewind_to(st.seq_len)?;
204        }
205        Ok(())
206    }
207}
metrale_model_arch/glm5next_mtp_head/proposer.rs

metrale_model_arch/glm5next_mtp_head/
proposer.rs
