cherenkov/qwen4_exp/gpu/
state.rs

1//! Commit/rollback, sequence reset, and prefix checkpoints.
2
3use super::*;
4
5impl Gpu<'_> {
6    /// Keep the first `n` rows of the step in flight: roll the recurrent
7    /// state back to after row n-1 if rows were rejected, advance the
8    /// position. KV rows and the PLE ring are positional and get
9    /// overwritten by the next step.
10    pub fn commit(&mut self, n: usize) -> Result<()> {
11        anyhow::ensure!(
12            n >= 1 && n <= self.batch_nb,
13            "commit {n} of {} rows",
14            self.batch_nb
15        );
16
17        if n < self.batch_nb {
18            anyhow::ensure!(self.batch_snap, "rollback needs a snapshotting step");
19
20            let plane = n - 1;
21            let cb = self.ctx.queue.commandBuffer().context("command buffer")?;
22            let enc = cb.computeCommandEncoder().context("encoder")?;
23
24            for l in &self.layers {
25                let Mix::Delta(d) = &l.mix else {
26                    continue;
27                };
28
29                for (src, dst) in [(&d.mid, &d.state), (&d.mid_hist, &d.hist)] {
30                    let len = (dst.length() / 4) as u32;
31
32                    self.dispatch(
33                        &enc,
34                        &self.pipes.copy_f32,
35                        |e| {
36                            self.bind(e, 0, src, plane * dst.length());
37                            self.bind(e, 1, dst, 0);
38                            set_bytes(e, 2, &len);
39                        },
40                        len as usize,
41                        256,
42                        false,
43                    );
44                }
45            }
46
47            enc.endEncoding();
48            cb.commit();
49            cb.waitUntilCompleted();
50        }
51
52        self.pos = self.batch_pos + n;
53
54        self.tokens.truncate(self.pos);
55
56        Ok(())
57    }
58
59    /// Forget the sequence: recurrent states, conv histories and the PLE
60    /// ring are zeroed; KV rows are overwritten positionally.
61    pub fn reset(&mut self) {
62        let zero = |b: &Buf| unsafe {
63            std::ptr::write_bytes(b.contents().cast::<u8>().as_ptr(), 0, b.length());
64        };
65
66        for l in &self.layers {
67            if let Mix::Delta(d) = &l.mix {
68                zero(&d.state);
69                zero(&d.hist);
70            }
71
72            if let Some(p) = &l.ple {
73                zero(&p.hist);
74            }
75        }
76
77        self.pos = 0;
78
79        self.tokens.clear();
80
81        self.batch_nb = 0;
82        self.mtp_len = 0;
83    }
84
85    pub(super) fn prefix_regions(&self, pos: usize, mtp_len: usize) -> Vec<(&Buf, usize, usize)> {
86        let c = &self.p.cfg;
87        let kv_row = c.num_key_value_heads * c.head_dim;
88        let scales = kv_q8_side(self.max_t, kv_row).0;
89        let mut regions = Vec::new();
90
91        for (l, n) in self
92            .layers
93            .iter()
94            .map(|l| (l, pos))
95            .chain(self.mtp.iter().map(|m| (&m.layer, mtp_len)))
96        {
97            match &l.mix {
98                Mix::Delta(d) => {
99                    regions.push((&d.state, 0, d.state.length()));
100                    regions.push((&d.hist, 0, d.hist.length()));
101                }
102                Mix::Attn(a) => {
103                    for b in [&a.kc, &a.vc] {
104                        regions.push((b, 0, n * kv_row));
105                        regions.push((b, scales, n * kv_row / 32 * 2));
106                    }
107
108                    regions.push((&a.ikc, 0, n * c.indexer_head_dim * 4));
109                    regions.push((
110                        &a.blk,
111                        0,
112                        n.div_ceil(c.indexer_compress_ratio) * c.indexer_head_dim * 4,
113                    ));
114                }
115            }
116
117            if let Some(ple) = &l.ple {
118                regions.push((&ple.hist, 0, ple.hist.length()));
119            }
120        }
121
122        regions
123    }
124
125    pub(crate) fn prefix_state_bytes(&self) -> usize {
126        self.state_bytes_at(self.pos, self.mtp_len)
127    }
128
129    pub(crate) fn state_bytes_at(&self, pos: usize, mtp_len: usize) -> usize {
130        self.prefix_regions(pos, mtp_len)
131            .iter()
132            .map(|(_, _, n)| n + size_of::<Vec<u8>>())
133            .sum::<usize>()
134            + pos * size_of::<u32>()
135            + size_of::<PrefixState>()
136    }
137
138    pub(crate) fn save_prefix(&self) -> PrefixState {
139        let data = self
140            .prefix_regions(self.pos, self.mtp_len)
141            .into_iter()
142            .map(|(b, off, n)| {
143                assert!(off + n <= b.length());
144
145                // Prefill has waited for all GPU work before checkpointing.
146                unsafe {
147                    std::slice::from_raw_parts(b.contents().cast::<u8>().as_ptr().add(off), n)
148                }
149                .to_vec()
150            })
151            .collect();
152
153        PrefixState {
154            pos: self.pos,
155            mtp_len: self.mtp_len,
156            tokens: self.tokens[..self.pos].to_vec(),
157            data,
158        }
159    }
160
161    /// Reuse suspended-request buffers; only initialized prefixes are copied.
162    pub(crate) fn save_into(&self, checkpoint: &mut Option<PrefixState>) {
163        let Some(state) = checkpoint else {
164            *checkpoint = Some(self.save_prefix());
165
166            return;
167        };
168
169        for ((buffer, offset, len), data) in self
170            .prefix_regions(self.pos, self.mtp_len)
171            .into_iter()
172            .zip(&mut state.data)
173        {
174            data.resize(len, 0);
175
176            unsafe {
177                std::ptr::copy_nonoverlapping(
178                    buffer.contents().cast::<u8>().as_ptr().add(offset),
179                    data.as_mut_ptr(),
180                    len,
181                );
182            }
183        }
184
185        state.pos = self.pos;
186        state.mtp_len = self.mtp_len;
187
188        state.tokens.clone_from(&self.tokens);
189    }
190
191    pub(crate) fn restore_prefix(&mut self, state: &PrefixState) -> Result<()> {
192        ensure!(
193            state.pos <= self.max_t && state.mtp_len <= self.max_t,
194            "prefix exceeds context"
195        );
196
197        let regions = self.prefix_regions(state.pos, state.mtp_len);
198
199        ensure!(
200            regions.len() == state.data.len(),
201            "prefix state layout mismatch"
202        );
203
204        for ((b, off, n), data) in regions.into_iter().zip(&state.data) {
205            ensure!(
206                n == data.len() && off + n <= b.length(),
207                "prefix region mismatch"
208            );
209
210            unsafe {
211                std::ptr::copy_nonoverlapping(
212                    data.as_ptr(),
213                    b.contents().cast::<u8>().as_ptr().add(off),
214                    n,
215                );
216            }
217        }
218
219        self.pos = state.pos;
220        self.mtp_len = state.mtp_len;
221
222        self.tokens.clone_from(&state.tokens);
223
224        self.batch_pos = self.pos;
225        self.batch_nb = 0;
226        self.folded_mtp = None;
227
228        Ok(())
229    }
230
231    /// Start another request without retaining unbounded profiling history.
232    /// Weights and expert residency remain loaded; sequence state is independent.
233    pub(crate) fn reset_request(&mut self) {
234        self.reset();
235        self.clear_profile();
236    }
237
238    /// Server scheduling publishes counters each turn instead of retaining a log.
239    pub(crate) fn clear_profile(&mut self) {
240        self.expert_history.clear();
241        self.route_history.clear();
242        self.state_history.clear();
243        self.ngram_ms.clear();
244        self.step_ms.clear();
245        self.gpu_ms.clear();
246        self.io_ms.clear();
247        self.misses.clear();
248        self.miss_bytes.clear();
249        self.warm.clear();
250        self.set_ms.clear();
251        self.read_ms.clear();
252        self.lookahead_hit.clear();
253        self.lookahead_issued.clear();
254        self.la_log.clear();
255        self.la_pending.clear();
256        self.cut.clear();
257        self.dispatches.clear();
258        self.gpu_idle_ms.clear();
259        self.rows.clear();
260        self.mtp_ms.clear();
261        self.prefill_stats.clear();
262    }
263}
cherenkov/qwen4_exp/gpu/state.rs

cherenkov/qwen4_exp/gpu/
state.rs
