metrale_model_engine/model/trait_impl/
impl_forward.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `impl ModelForward for TransformerModel`, mostly delegating to `<method>_dispatch`
4//! helpers in the sibling modules.
5//!
6//! Owner: model-engine.
7//! Invariants: the ones in `trait_impl/mod.rs`.
8
9use anyhow::Result;
10use metrale_gpu_runtime::gpu::DevicePtr;
11
12use crate::model::types::TransformerModel;
13use crate::traits::{ModelForward, PrefillSlice, SequenceState};
14
15impl ModelForward for TransformerModel {
16    // 2026-09-25: `prefill`, `prefill_chunk`, `prefill_twophase` and
17    // `mixed_forward` each call `try_eager_drafter_prefill` after the forward.
18    // The whole-prompt drafter capture (`mtp_prefill_hidden`) is one slot shared
19    // by all sequences, and a first chunk of any sequence restarts it, so it is
20    // consumed while this sequence still owns it. `METRALE_NO_MTP_EAGER_DRAFTER`
21    // turns the eager consume off.
22    fn prefill(&self, tokens: &[u32], seq: &mut SequenceState, stream: u64) -> Result<DevicePtr> {
23        self.gdn_carry_flush_pending()?;
24        self.stamp_overlay_route(seq.adapter_slot);
25        let logits = self.prefill_dispatch(tokens, seq, stream)?;
26        self.try_eager_drafter_prefill(seq, true, stream);
27        Ok(logits)
28    }
29
30    fn prefill_chunk(
31        &self,
32        tokens: &[u32],
33        seq: &mut SequenceState,
34        chunk_start: usize,
35        chunk_len: usize,
36        is_last_chunk: bool,
37        stream: u64,
38    ) -> Result<DevicePtr> {
39        self.gdn_carry_flush_pending()?;
40        self.stamp_overlay_route(seq.adapter_slot);
41        let logits = self.prefill_chunk_dispatch(
42            tokens,
43            seq,
44            chunk_start,
45            chunk_len,
46            is_last_chunk,
47            stream,
48        )?;
49        self.try_eager_drafter_prefill(seq, is_last_chunk, stream);
50        Ok(logits)
51    }
52
53    fn prefill_twophase(
54        &self,
55        tokens: &[u32],
56        seq: &mut SequenceState,
57        chunk_size: usize,
58        stream: u64,
59    ) -> Result<DevicePtr> {
60        self.gdn_carry_flush_pending()?;
61        self.stamp_overlay_route(seq.adapter_slot);
62        let logits = self.prefill_twophase_dispatch(tokens, seq, chunk_size, stream)?;
63        self.try_eager_drafter_prefill(seq, true, stream);
64        Ok(logits)
65    }
66
67    fn decode(&self, token: u32, seq: &mut SequenceState, _stream: u64) -> Result<DevicePtr> {
68        self.gdn_carry_flush_pending()?;
69        self.stamp_overlay_route(seq.adapter_slot);
70        self.stamp_decode_moe_single(seq.adapter_slot);
71        self.decode_dispatch(token, seq, _stream)
72    }
73
74    fn decode_batch(
75        &self,
76        tokens: &[u32],
77        seqs: &mut [&mut SequenceState],
78        stream: u64,
79    ) -> Result<DevicePtr> {
80        self.gdn_carry_flush_pending()?;
81        self.stamp_overlay_route_batch(seqs);
82        self.stamp_decode_moe_batch(seqs);
83        let r = self.decode_batch_dispatch(tokens, seqs, stream);
84        if r.is_err() {
85            // 2026-09-25: An error raised during graph capture (for example
86            // `reject_decode_lora` refusing an adapter-routed MoE row) leaves the
87            // capture open. End and discard it so the caller's cleanup does not
88            // run on a capturing stream. The batched decode captures on the
89            // default stream (decode_a2.rs).
90            self.gpu.abort_capture_if_active(self.gpu.default_stream());
91        }
92        r
93    }
94
95    fn mixed_forward(
96        &self,
97        decode_tokens: &[u32],
98        decode_seqs: &mut [&mut SequenceState],
99        prefill_tokens: &[u32],
100        prefill_seq: &mut SequenceState,
101        prefill_chunk_start: usize,
102        prefill_chunk_len: usize,
103        prefill_is_last: bool,
104        stream: u64,
105    ) -> Result<crate::traits::MixedForwardResult> {
106        self.gdn_carry_flush_pending()?;
107        // 2026-09-25: A mixed step always stamps `i32::MIN` (mixed adapters), so
108        // the token-overlay hooks skip for the whole step.
109        self.overlay_route_slot
110            .store(i32::MIN, std::sync::atomic::Ordering::Relaxed);
111        self.stamp_decode_moe_batch(decode_seqs);
112        let r = self.mixed_forward_dispatch(
113            decode_tokens,
114            decode_seqs,
115            prefill_tokens,
116            prefill_seq,
117            prefill_chunk_start,
118            prefill_chunk_len,
119            prefill_is_last,
120            stream,
121        );
122        if r.is_err() {
123            // 2026-09-25: The same capture release as in `decode_batch`.
124            self.gpu.abort_capture_if_active(self.gpu.default_stream());
125        }
126        let out = r?;
127        self.try_eager_drafter_prefill(prefill_seq, prefill_is_last, stream);
128        Ok(out)
129    }
130
131    /// 2026-09-25: `prefill_batch_chunk_dispatch` sends a batch that is
132    /// ineligible, or whose cache plan is not admitted, to its per-stream loop
133    /// before any sequence is mutated. An error from an admitted kernel batch is
134    /// returned: a per-stream retry could apply prefix-cache and KV state twice.
135    fn prefill_batch_chunk(
136        &self,
137        streams: &mut [PrefillSlice<'_>],
138        stream: u64,
139    ) -> Result<Vec<DevicePtr>> {
140        self.gdn_carry_flush_pending()?;
141        self.prefill_batch_chunk_rows(streams, stream, 0)
142    }
143
144    /// 2026-09-25: Mixed-step variant: finishing streams write their logits from
145    /// row `row_base`, clear of the decode lanes. The trait docs describe the
146    /// aliasing this prevents.
147    fn prefill_batch_chunk_rows(
148        &self,
149        streams: &mut [PrefillSlice<'_>],
150        stream: u64,
151        row_base: usize,
152    ) -> Result<Vec<DevicePtr>> {
153        self.gdn_carry_flush_pending()?;
154        self.prefill_batch_chunk_dispatch(streams, stream, row_base)
155    }
156
157    fn normalize_ssm_states(&self, seq: &SequenceState, stream: u64) -> Result<()> {
158        self.gdn_carry_flush_pending()?;
159        self.normalize_ssm_states_dispatch(seq, stream)
160    }
161
162    fn hc_mult(&self) -> usize {
163        self.config.hc_mult
164    }
165
166    fn is_mla(&self) -> bool {
167        self.is_mla_dispatch()
168    }
169
170    fn kv_block_size(&self) -> Option<usize> {
171        Some(self.kv_cache.lock().block_size())
172    }
173
174    fn prefill_tail_split(&self, tokens: &[u32]) -> Option<usize> {
175        self.prefill_tail_split_dispatch(tokens)
176    }
177}
metrale_model_engine/model/trait_impl/impl_forward.rs

metrale_model_engine/model/trait_impl/
impl_forward.rs
