metrale_model_layers/layers/
mtp_multi.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Multi-module MTP proposer: draft slot i runs its own
4//! `MtpHead` (own weights and KV cache), as
5//! `modules[i].forward_one(previous draft token, previous module's hidden)`.
6//!
7//! Owner: model-layers (MTP head).
8//! Invariants:
9//! - `modules` is non-empty, and each state holds one `MtpProposerState`
10//!   per module.
11//! - A propose runs `min(num_drafts, modules.len())` modules, the first ones.
12
13use std::any::Any;
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
17
18use crate::layer::ForwardContext;
19use crate::layers::mtp_head::{MtpHead, MtpProposerState};
20use crate::speculative::{DraftProposer, ProposerState};
21
22/// 2026-09-25: Per-sequence state for `MultiModuleMtpHead`: one
23/// `MtpProposerState` per module, since the modules do not share a KV cache.
24pub struct MultiModuleMtpState {
25    /// 2026-09-25: `per_module[i]` belongs to `modules[i]`; `alloc_state`
26    /// builds one per module.
27    pub per_module: Vec<MtpProposerState>,
28    /// 2026-09-25: Drafts produced by the last `propose`; `after_verify`
29    /// updates that many modules.
30    pub last_num_drafted: usize,
31}
32
33impl ProposerState for MultiModuleMtpState {
34    fn as_any(&self) -> &dyn Any {
35        self
36    }
37    fn as_any_mut(&mut self) -> &mut dyn Any {
38        self
39    }
40}
41
42/// 2026-09-25: Independent MTP modules, one per draft slot.
43pub struct MultiModuleMtpHead {
44    /// 2026-09-25: Non-empty (`new` refuses an empty list). `build_mtp_proposer`
45    /// builds one head per loaded MTP weight set.
46    modules: Vec<MtpHead>,
47}
48
49impl std::fmt::Debug for MultiModuleMtpHead {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        f.debug_struct("MultiModuleMtpHead")
52            .field("modules", &self.modules.len())
53            .finish()
54    }
55}
56
57impl MultiModuleMtpHead {
58    /// 2026-09-25: Assemble a multi-module proposer from per-module heads.
59    /// Errors on an empty list.
60    pub fn new(modules: Vec<MtpHead>) -> Result<Self> {
61        anyhow::ensure!(
62            !modules.is_empty(),
63            "MultiModuleMtpHead requires at least one module (got 0); \
64             caller should not construct this type for single-module MTP"
65        );
66        Ok(Self { modules })
67    }
68
69    /// 2026-09-25: Number of modules; `propose` caps `num_drafts` at it.
70    pub fn num_modules(&self) -> usize {
71        self.modules.len()
72    }
73}
74
75impl DraftProposer for MultiModuleMtpHead {
76    fn alloc_state(&self, _gpu: &dyn GpuBackend) -> Result<Box<dyn ProposerState>> {
77        let per_module = (0..self.modules.len())
78            .map(|_| MtpProposerState {
79                block_table: Vec::new(),
80                seq_len: 0,
81                last_num_drafted: 0,
82                last_pair_key: None,
83                last_drafts: Vec::new(),
84                pending_catchup: Vec::new(),
85            })
86            .collect();
87        Ok(Box::new(MultiModuleMtpState {
88            per_module,
89            last_num_drafted: 0,
90        }))
91    }
92
93    fn propose(
94        &self,
95        last_token: u32,
96        target_hidden: DevicePtr,
97        position: usize,
98        num_drafts: usize,
99        state: &mut dyn ProposerState,
100        ctx: &ForwardContext,
101        stream: u64,
102        draft_embed_target: Option<DevicePtr>,
103        grammar_bitmask: Option<&[i32]>,
104        _target_hidden_stack: Option<DevicePtr>,
105    ) -> Result<Vec<u32>> {
106        let mm_state = state
107            .as_any_mut()
108            .downcast_mut::<MultiModuleMtpState>()
109            .ok_or_else(|| anyhow::anyhow!("Invalid MultiModuleMtp state"))?;
110
111        // 2026-09-25: More drafts than modules are clamped to the module
112        // count without a warning.
113        let k = num_drafts.min(self.modules.len());
114
115        let mut drafts = Vec::with_capacity(k);
116        let mut current_token = last_token;
117        let mut current_hidden = target_hidden;
118
119        for i in 0..k {
120            // 2026-09-25: Only the last draft stages its embedding at
121            // `draft_embed_target`; earlier drafts reach the next module as
122            // the returned token.
123            let embed_target = if i == k - 1 { draft_embed_target } else { None };
124
125            // 2026-09-25: Every module gets the same single-position grammar
126            // mask.
127            let mask_for_draft = grammar_bitmask;
128
129            let draft = self.modules[i].forward_one(
130                current_token,
131                current_hidden,
132                position + i,
133                &mut mm_state.per_module[i],
134                ctx,
135                stream,
136                embed_target,
137                mask_for_draft,
138                i == 0,
139            )?;
140
141            tracing::debug!(
142                "MultiMTP propose[{i}/{k}]: token={current_token} pos={} module_seq_len={} → draft={draft}",
143                position + i,
144                mm_state.per_module[i].seq_len,
145            );
146
147            drafts.push(draft);
148            current_token = draft;
149            // 2026-09-25: Module i + 1 reads module i's residual stream from
150            // `hidden_states`, whatever `ModelLevers::mtp_chain_postnorm` says
151            // (`MtpHead::chain_hidden` honours it).
152            current_hidden = ctx.buffers.hidden_states();
153        }
154
155        mm_state.last_num_drafted = drafts.len();
156        Ok(drafts)
157    }
158
159    fn read_deferred_draft_token(&self, gpu: &dyn GpuBackend) -> Result<u32> {
160        // 2026-09-25: Reads the last module's deferred id. The last draft
161        // came from `modules[k - 1]`, so this is that draft only when the
162        // propose ran every module (k == modules.len()).
163        self.modules
164            .last()
165            .expect("MultiModuleMtpHead::new enforces non-empty")
166            .read_deferred_draft_token(gpu)
167    }
168
169    fn after_verify(
170        &self,
171        num_accepted: usize,
172        state: &mut dyn ProposerState,
173        stream: u64,
174    ) -> Result<()> {
175        let mm_state = state
176            .as_any_mut()
177            .downcast_mut::<MultiModuleMtpState>()
178            .ok_or_else(|| anyhow::anyhow!("Invalid MultiModuleMtp state"))?;
179
180        // 2026-09-25: Each of the k modules wrote one row in the last
181        // propose. Modules 0..num_accepted keep it; modules num_accepted..k
182        // drop it.
183        let k = mm_state.last_num_drafted;
184        for (i, per) in mm_state.per_module.iter_mut().take(k).enumerate() {
185            per.last_num_drafted = 1;
186            let trim = if i < num_accepted { 0 } else { 1 };
187            if trim > 0 {
188                per.seq_len = per.seq_len.saturating_sub(trim);
189            }
190        }
191        let _ = stream;
192        tracing::debug!(
193            "MultiMTP after_verify: accepted={num_accepted} of {k}; per-module trim done"
194        );
195        Ok(())
196    }
197
198    fn free_state(&self, _gpu: &dyn GpuBackend, state: &mut dyn ProposerState) -> Result<()> {
199        let mm_state = state
200            .as_any_mut()
201            .downcast_mut::<MultiModuleMtpState>()
202            .ok_or_else(|| anyhow::anyhow!("Invalid MultiModuleMtp state"))?;
203        // 2026-09-25: Return each module's blocks to that module's KV cache.
204        for (i, per) in mm_state.per_module.iter_mut().enumerate() {
205            let head = &self.modules[i];
206            if !per.block_table.is_empty() {
207                head.kv_cache_lock().free_blocks(&per.block_table);
208                per.block_table.clear();
209            }
210            per.seq_len = 0;
211        }
212        Ok(())
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[test]
221    fn test_empty_modules_rejected() {
222        let err = MultiModuleMtpHead::new(vec![]).unwrap_err();
223        assert!(
224            err.to_string().contains("at least one module"),
225            "unexpected error: {err}"
226        );
227    }
228}
metrale_model_layers/layers/mtp_multi.rs

metrale_model_layers/layers/
mtp_multi.rs
