spark_model/layers/deepseek_v41_layer/
step_seg.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! The single-token step as a few CUDA graphs spanning layers
5//! (`ATLAS_DS41_STEP_GRAPH=1`, needs `ATLAS_DS41_DEVICE_ROUTE=1`).
6//!
7//! With the routing on the device and every engram layer's rows uploaded
8//! up front, the only host work left in a decode step is the index-source
9//! layers' attention (the index score read-back and the host top-k, plus
10//! the compressor's position parity): eight layers of forty. So the step is
11//! captured as SEGMENTS between those points: segment 0 = layers 0 and 1
12//! through layer 2's attention input; segment k = the ffn of index layer k
13//! through the next index layer's attention input; the last runs to the
14//! final collapse. Nine graph launches a token instead of ~1,900 kernel
15//! launches, and one read-back per segment (the routing headers of its
16//! layers) instead of one per layer.
17//!
18//! The protocol. A segment's graph carries a SNAPSHOT before every layer's
19//! ffn (streams, the attention output, pre_a, the attention site's post and
20//! comb mixes: ~75 KB a layer, device copies inside the graph), and a segment is launched only
21//! at its OWNER's step, never at its close. The owner reads its layers'
22//! routing headers once after the launch. Inside a replay nothing fetches:
23//! the selection kernel points an absent expert at slot 0 and flags the
24//! layer, and every layer after it routed on a wrong input. The first
25//! flagged layer k has real picks (its input was right): its misses are
26//! fetched, its snapshot restored, layer k runs from its ffn on with the
27//! eager step's own code (`ffn_eager`), the layers after it to the
28//! segment's end run the whole eager step, and layers s..k-1 stand. The
29//! CAPTURE token records the bodies on a second stream (the capture is
30//! relaxed) while every layer runs eagerly on the compute stream, so the
31//! token that captures is right by the eager path and the first replay is
32//! the next token. An owner does its own attention input eagerly whenever
33//! the segment before it did not replay through to it. A failed capture
34//! turns the mode off for the process; the token in flight finishes
35//! eagerly.
36//!
37//! Every layer's `step` lands here; the OWNER of a segment (layer 0, or an
38//! index layer at its ffn) prepares, launches and runs the protocol, and
39//! the layers a replayed segment covered return at once (`skip_until`).
40
41use std::sync::OnceLock;
42
43use anyhow::{Context, Result, ensure};
44use spark_runtime::gpu::{DevicePtr, GpuBackend, GraphHandle};
45
46use super::step::{trace_hash, trace_on};
47use super::step_seg_run::SAVE_N;
48use super::{DeepSeekV41Layer, V41LayerState};
49use crate::layer::ForwardContext;
50use crate::layers::attn_v41::LayerRole;
51use crate::layers::moe_v41::MoeV41;
52
53/// `ATLAS_DS41_STEP_GRAPH=1`, read once.
54pub fn step_graph_on() -> bool {
55    static V: OnceLock<bool> = OnceLock::new();
56    *V.get_or_init(|| std::env::var("ATLAS_DS41_STEP_GRAPH").is_ok_and(|v| v == "1"))
57}
58
59/// One captured segment.
60pub struct SegGraph {
61    pub(super) g: GraphHandle,
62    /// The layer whose attention input closes the segment (`n_layers` for
63    /// the last one); layers below it are covered.
64    pub(super) end: usize,
65    /// Layers whose MoE ran inside (their headers to read after a launch).
66    pub(super) moe_layers: Vec<u32>,
67    /// Selection classes the covered attention bodies read (ratio > 0,
68    /// ratio 0) and the compressed rows they bake.
69    pub(super) classes: (bool, bool),
70    pub(super) rows_b: Option<DevicePtr>,
71}
72
73/// An open capture: the owner and the MoE layers appended so far.
74pub struct SegCapture {
75    pub(super) owner: usize,
76    pub(super) moe_layers: Vec<u32>,
77    pub(super) classes: (bool, bool),
78    pub(super) rows_b: Option<DevicePtr>,
79}
80
81#[derive(Default)]
82pub struct SegState {
83    /// A capture failed: the mode is off for the process.
84    pub disabled: bool,
85    /// Layers below this were covered by the segment just launched.
86    pub skip_until: usize,
87    /// Layers below this run the whole eager step this token.
88    pub eager_until: usize,
89    /// The layer whose snapshot the owner restored: it runs from its ffn on.
90    pub ffn_only_at: Option<usize>,
91    pub capturing: Option<SegCapture>,
92    /// By owner layer.
93    pub graphs: Vec<Option<SegGraph>>,
94    /// The buffers every graph bakes: hidden, streams, normed.
95    pub baked: Option<[DevicePtr; 3]>,
96    /// The stream a capture token records on, the side stream the recorded
97    /// shared expert forks onto, and the fork / join events (all captured
98    /// into the graph; the eager path's own side stream is never touched).
99    pub cap_stream: Option<u64>,
100    pub cap_side: u64,
101    pub ev_fork: u64,
102    pub ev_join: u64,
103    /// Per layer, the snapshot before its ffn (`step_seg_run::SAVE_N`
104    /// buffers: hidden, streams, pre_prev, pre_a, the attention output,
105    /// post_s, comb_s).
106    pub save: Vec<[DevicePtr; SAVE_N]>,
107    /// Every layer's window ring as last seen (`note_window`): a sequence's
108    /// own buffer the captured attention bodies bake.
109    pub windows: Vec<DevicePtr>,
110    /// A window ring changed since the graphs were captured: recapture.
111    pub stale: bool,
112    /// Graph launches and miss falls this process, for the log.
113    pub launches: u64,
114    pub reruns: u64,
115}
116
117impl DeepSeekV41Layer {
118    fn role_of(&self, l: usize) -> Result<LayerRole> {
119        self.rt.roles.lock().unwrap()[l].context("step graph: a layer without a role")
120    }
121
122    fn capturable(r: &LayerRole) -> bool {
123        !r.is_kv_source && !r.is_index_source
124    }
125
126    /// The segment an owner at `s` spans: its end and the selection classes
127    /// of the capturable layers in it.
128    fn geometry(&self, s: usize) -> Result<(usize, (bool, bool))> {
129        let n = self.rt.n_layers;
130        let mut classes = (false, false);
131        let mut end = n;
132        for l in s..n {
133            let r = self.role_of(l)?;
134            if l != s && !Self::capturable(&r) {
135                end = l;
136                break;
137            }
138            if Self::capturable(&r) {
139                if r.ratio > 0 {
140                    classes.0 = true;
141                } else {
142                    classes.1 = true;
143                }
144            }
145        }
146        Ok((end, classes))
147    }
148
149    /// Every step (any mode, prefill too) notes this layer's window ring: the
150    /// captured attention bodies bake it, and a new sequence brings new
151    /// rings, so a change marks the graphs stale.
152    pub(super) fn note_window(&self, window: DevicePtr) {
153        let mut seg = self.rt.seg.lock().unwrap();
154        let n = self.rt.n_layers;
155        if seg.windows.len() != n {
156            seg.windows = vec![DevicePtr(0); n];
157        }
158        if seg.windows[self.idx] != window {
159            seg.windows[self.idx] = window;
160            seg.stale = true;
161        }
162    }
163
164    /// The step of one layer in segment-graph mode.
165    pub(super) fn step_seg(
166        &self,
167        hidden: DevicePtr,
168        start_pos: usize,
169        st: &mut V41LayerState,
170        ctx: &ForwardContext,
171        stream: u64,
172    ) -> Result<()> {
173        let rt = &self.rt;
174        let gpu = ctx.gpu;
175        let streams = ctx.buffers.hc_streams();
176        let normed = ctx.buffers.norm_output();
177        let n = rt.n_layers;
178        let mut seg = rt.seg.lock().unwrap();
179        if seg.graphs.len() != n {
180            seg.graphs = (0..n).map(|_| None).collect();
181        }
182        if self.idx == 0 {
183            // a new step: the baked buffers, the skip and fallback marks
184            seg.skip_until = 0;
185            seg.eager_until = 0;
186            seg.ffn_only_at = None;
187            let baked = [hidden, streams, normed];
188            if seg.baked != Some(baked) || seg.stale {
189                if seg.baked.is_some() && seg.graphs.iter().any(Option::is_some) {
190                    tracing::info!(
191                        "DS41 step graph: {}; recapturing every segment",
192                        if seg.stale {
193                            "a sequence's window rings changed"
194                        } else {
195                            "buffers changed"
196                        }
197                    );
198                }
199                for g in seg.graphs.iter_mut().flat_map(Option::take) {
200                    gpu.destroy_graph(g.g)?;
201                }
202                seg.baked = Some(baked);
203                seg.stale = false;
204            }
205            if seg.cap_stream.is_none() {
206                seg.cap_stream = Some(
207                    gpu.create_stream()
208                        .context("step graph: the recording stream")?,
209                );
210                seg.cap_side = gpu.create_stream().context("step graph: the side stream")?;
211                seg.ev_fork = gpu.create_event()?;
212                seg.ev_join = gpu.create_event()?;
213            }
214            if seg.save.len() != n {
215                self.alloc_save(gpu, &mut seg)?;
216            }
217        }
218        if self.idx < seg.skip_until {
219            return Ok(());
220        }
221        if seg.ffn_only_at == Some(self.idx) {
222            // the owner restored this layer's snapshot: from the ffn on
223            seg.ffn_only_at = None;
224            drop(seg);
225            return self.ffn_eager(hidden, 1, start_pos, ctx, stream);
226        }
227        if self.idx < seg.eager_until {
228            drop(seg);
229            return self.eager_layer(hidden, start_pos, st, ctx, stream);
230        }
231        let role = self.role_of(self.idx)?;
232        let owner = self.idx == 0 || !Self::capturable(&role);
233        let cs = seg.cap_stream.context("step graph: no recording stream")?;
234
235        // a layer inside an open capture: its body is recorded on the
236        // recording stream while the layer runs eagerly on the compute stream
237        if let Some(mut cap) = seg.capturing.take() {
238            ensure!(
239                cap.owner < self.idx,
240                "step graph: capture owned by a later layer"
241            );
242            if owner {
243                // the closer: its attention input ends the segment; it runs
244                // that input eagerly below, as any owner a replay did not reach
245                let r = (|| -> Result<()> {
246                    self.engram_apply(gpu, streams, cs)?;
247                    self.seg_attn_in(gpu, streams, hidden, normed, cs)
248                })();
249                self.close_or_disable(gpu, &mut seg, cap, self.idx, cs, r);
250            } else {
251                let arena = MoeV41::arena_of(&rt.lru.lock().unwrap());
252                let save = seg.save[self.idx];
253                let fork = (seg.cap_side, seg.ev_fork, seg.ev_join);
254                let r = (|| -> Result<()> {
255                    let moe = rt.moe.lock().unwrap();
256                    self.body_attn(gpu, st, hidden, streams, normed, cap.rows_b, cs)?;
257                    self.save_nodes(gpu, &save, hidden, streams, cs)?;
258                    self.body_ffn(gpu, &moe, arena, fork, hidden, streams, normed, cs)
259                })();
260                cap.moe_layers.push(self.idx as u32);
261                if r.is_err() || self.idx + 1 == n {
262                    self.close_or_disable(gpu, &mut seg, cap, n, cs, r);
263                } else {
264                    seg.capturing = Some(cap);
265                }
266                drop(seg);
267                return self.eager_layer(hidden, start_pos, st, ctx, stream);
268            }
269        }
270        if seg.disabled {
271            drop(seg);
272            return self.eager_layer(hidden, start_pos, st, ctx, stream);
273        }
274        ensure!(
275            owner,
276            "step graph: layer {} reached outside any segment",
277            self.idx
278        );
279
280        // the owner: its attention input if no replay delivered it, then
281        // prepare, then replay or capture the segment from here
282        if self.idx > 0 && seg.skip_until != self.idx {
283            self.engram_apply(gpu, streams, stream)?;
284            self.seg_attn_in(gpu, streams, hidden, normed, stream)?;
285        }
286        let (end, classes) = self.geometry(self.idx)?;
287        let rows_b = self.prepare(gpu, st, ctx, start_pos, classes, normed, stream)?;
288        if let Some(sg) = seg.graphs[self.idx].as_ref()
289            && (sg.end != end || sg.classes != classes || sg.rows_b != rows_b)
290        {
291            tracing::warn!(
292                "DS41 step graph: segment {} bakes other rows; recapturing",
293                self.idx
294            );
295            if let Some(sg) = seg.graphs[self.idx].take() {
296                gpu.destroy_graph(sg.g)?;
297            }
298        }
299        if let Some(sg) = seg.graphs[self.idx].as_ref() {
300            let (gh, layers) = (sg.g, sg.moe_layers.clone());
301            let first_miss = self.launch_once(gpu, gh, &layers, stream)?;
302            seg.launches += 1;
303            let Some(k) = first_miss else {
304                seg.skip_until = end;
305                if trace_on() {
306                    // `ATLAS_DS41_TRACE=1`: the highway after the segment, to
307                    // diff against the eager trace's line for layer end-1
308                    gpu.synchronize(stream)?;
309                    tracing::info!(
310                        "DS41 trace pos={start_pos} seg={}..{end} hwy={:016x} hidden={:016x}",
311                        self.idx,
312                        trace_hash(gpu, streams, rt.hc_mult * rt.hidden * 4),
313                        trace_hash(gpu, hidden, rt.hidden * 2)
314                    );
315                }
316                if end == n {
317                    self.log_step(&seg, start_pos);
318                }
319                return Ok(());
320            };
321            // layer k's picks were real and are resident now: its snapshot
322            // back, k from its ffn on and k+1..end whole, eagerly
323            seg.reruns += 1;
324            let save = seg.save[k];
325            self.restore(gpu, &save, hidden, streams, stream)?;
326            seg.eager_until = end;
327            if k == self.idx {
328                drop(seg);
329                return self.ffn_eager(hidden, 1, start_pos, ctx, stream);
330            }
331            seg.skip_until = k;
332            seg.ffn_only_at = Some(k);
333            return Ok(());
334        }
335        // no graph yet: record from here on the recording stream while this
336        // token runs eagerly on the compute stream
337        let arena = MoeV41::arena_of(&rt.lru.lock().unwrap());
338        let save = seg.save[self.idx];
339        let fork = (seg.cap_side, seg.ev_fork, seg.ev_join);
340        let r = gpu
341            .begin_capture(cs)
342            .context("step graph: begin_capture")
343            .and_then(|_| {
344                let moe = rt.moe.lock().unwrap();
345                if self.idx == 0 {
346                    self.body_attn(gpu, st, hidden, streams, normed, rows_b, cs)?;
347                }
348                self.save_nodes(gpu, &save, hidden, streams, cs)?;
349                self.body_ffn(gpu, &moe, arena, fork, hidden, streams, normed, cs)
350            });
351        let cap = SegCapture {
352            owner: self.idx,
353            moe_layers: vec![self.idx as u32],
354            classes,
355            rows_b,
356        };
357        if r.is_err() || self.idx + 1 == n {
358            self.close_or_disable(gpu, &mut seg, cap, n, cs, r);
359        } else {
360            seg.capturing = Some(cap);
361        }
362        drop(seg);
363        if self.idx == 0 {
364            self.eager_layer(hidden, start_pos, st, ctx, stream)
365        } else {
366            // an index owner: its attention ran eagerly in `prepare`
367            self.ffn_eager(hidden, 1, start_pos, ctx, stream)
368        }
369    }
370
371    /// End the open capture as the segment of `cap.owner` ending at `end`,
372    /// or, when the recording failed (`r`) or the capture cannot end, turn
373    /// the mode off: the token in flight finishes eagerly.
374    fn close_or_disable(
375        &self,
376        gpu: &dyn GpuBackend,
377        seg: &mut SegState,
378        cap: SegCapture,
379        end: usize,
380        cs: u64,
381        r: Result<()>,
382    ) {
383        match r.and_then(|_| gpu.end_capture(cs).context("step graph: end_capture")) {
384            Ok(g) => {
385                tracing::info!(
386                    "DS41 step graph: segment {} = layers {}..{} captured ({} MoE layers)",
387                    cap.owner,
388                    cap.owner,
389                    end,
390                    cap.moe_layers.len()
391                );
392                seg.graphs[cap.owner] = Some(SegGraph {
393                    g,
394                    end,
395                    moe_layers: cap.moe_layers,
396                    classes: cap.classes,
397                    rows_b: cap.rows_b,
398                });
399            }
400            Err(e) => {
401                gpu.abort_capture_if_active(cs);
402                seg.disabled = true;
403                seg.capturing = None;
404                seg.eager_until = self.rt.n_layers;
405                tracing::warn!(
406                    "DS41 L{}: step graph capture failed ({e:#}); the mode is off, the step runs eagerly",
407                    self.idx
408                );
409            }
410        }
411    }
412
413    fn log_step(&self, seg: &SegState, start_pos: usize) {
414        if step_log_on() {
415            tracing::info!(
416                "DS41 step graph pos {start_pos}: {} launches, {} miss falls so far",
417                seg.launches,
418                seg.reruns
419            );
420        }
421    }
422}
423
424/// `ATLAS_DS41_STEP_GRAPH_LOG=1`: one line per token with the running counts.
425pub(super) fn step_log_on() -> bool {
426    static V: OnceLock<bool> = OnceLock::new();
427    *V.get_or_init(|| std::env::var("ATLAS_DS41_STEP_GRAPH_LOG").is_ok_and(|v| v == "1"))
428}
spark_model/layers/deepseek_v41_layer/step_seg.rs

spark_model/layers/deepseek_v41_layer/
step_seg.rs
