memra_engine/
prime_graph.rs

1//! PrimeGraph (task #14, design v3): a per-bucket CUDA graph of the FULL fresh-prime
2//! trunk, bound to a dedicated SCRATCH cache; serving replays it (one cuGraphLaunch,
3//! ~23ms vs ~26ms eager at bucket 512) and COPIES the outputs into the session's cache
4//! (KV rows + conv rings + recurrent states, ~tens of us D2D — the copy-out beats both
5//! table-indirect kernels and graphExec node patching, ledger design v3).
6//!
7//! Correctness story (all bit-proven by prime-graph-smoke + the gate):
8//! - fresh-prime semantics are BAKED as graph-head memset nodes (state/ring/len_d zero);
9//! - pads past the true length are invisible (gdn_pad_mask identity steps, causal
10//!   attention, device-indexed last-row gathers) — replay logits are bit-identical to
11//!   the eager true-length prime;
12//! - the GRAPH-OUTPUT CONTRACT: only the stable IO buffers and the scratch cache's
13//!   resident state survive a launch (in-graph transient addresses recycle).
14//! - ssm ping-pong: the capture-time core swapped the scratch cache's host fields; after
15//!   capture they name exactly the buffer the graph WRITES, and no further swaps happen,
16//!   so `scratch.recur[il].ssm_state` is the copy-out source on every replay.
17
18use crate::Engine;
19use crate::cache::Cache;
20use crate::hybrid::HybridModel;
21use cudarc::driver::{CudaGraph, CudaSlice};
22
23pub struct PrimeGraph {
24    pub bucket: usize,
25    graph: CudaGraph,
26    /// CAPTURE-RETAIN keeper (draft-graph law): holds every allocation the closure made so
27    /// the pool NEVER re-issues the graph's baked addresses to later eager work — without
28    /// it, any post-capture allocation can land on graph-internal addresses and every
29    /// replay scribbles it (the prime-graph-gate T=512 corruption, 2026-07-26).
30    _keeper: Vec<Box<dyn std::any::Any + Send>>,
31    /// PRIVATE f16 scratch (defect-hunt lead, 2026-07-26): the graph bakes the resident
32    /// f16 scratch's cvt/Lt pointers; sharing them with eager GEMMs cross-contaminates
33    /// replays. This scratch was resident DURING capture and is swapped back in around
34    /// every replay so the baked pointers always address graph-owned memory.
35    // held for lifetime only: resident during capture, swapped around replays
36    #[allow(dead_code)]
37    private_scratch: Option<crate::f16_ffi::F16Scratch>,
38    scratch: Cache,
39    x_in: CudaSlice<f32>,
40    len_d: CudaSlice<i32>,
41    logits_out: CudaSlice<f32>,
42    h_seed_out: CudaSlice<f32>,
43    n_embd: usize,
44}
45
46impl PrimeGraph {
47    /// Gate/debug accessor: the graph's bound scratch cache (read-only).
48    pub fn scratch(&self) -> &Cache {
49        &self.scratch
50    }
51}
52
53impl HybridModel {
54    /// Capture the fresh-prime graph for `bucket` tokens (13-15ms measured). Manual staged
55    /// capture — capture_graph_retained's keeper path trips on the prime (smoke finding 4).
56    pub fn prime_graph_new(
57        &self,
58        e: &Engine,
59        bucket: usize,
60    ) -> Result<PrimeGraph, Box<dyn std::error::Error>> {
61        self.refuse_hyper("prime_graph_new")?;
62        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
63        let n_embd = self.cfg.n_embd as usize;
64        let n_vocab = self.output.out_features();
65        let mut scratch = Cache::new(e, &self.cfg, bucket + 8)?;
66        let x_in = e.zeros(bucket * n_embd)?;
67        let pos_d = e.htod_i32(&(0..bucket as i32).collect::<Vec<_>>())?;
68        let len_d = e.htod_i32(&[bucket as i32])?;
69        let mut logits_out = e.uninit(n_vocab)?;
70        let mut h_seed_out = e.uninit(n_embd)?;
71
72        // capture with a PRIVATE f16 scratch resident (pre-sized to the trunk's largest
73        // GEMM input: m = bucket, in_f up to n_ff) so no eager call ever mutates the
74        // graph-baked buffers.
75        let n_ff_max = self
76            .layers
77            .iter()
78            .map(|l| match &l.ffn {
79                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
80                _ => n_embd,
81            })
82            .max()
83            .unwrap_or(n_embd)
84            .max(n_embd);
85        let private = crate::f16_ffi::F16Scratch::with_capacity(e, bucket * n_ff_max * 2)?;
86        let prev_scratch = e.f16_scratch_swap(Some(private));
87        let scratch_cell = std::cell::RefCell::new(&mut scratch);
88        let lo_cell = std::cell::RefCell::new(&mut logits_out);
89        let hs_cell = std::cell::RefCell::new(&mut h_seed_out);
90        let (graph, keeper) = e.capture_graph_retained(|e| {
91            let sc: &mut Cache = &mut scratch_cell.borrow_mut();
92            for kvl in sc.kv.iter_mut().flatten() {
93                kvl.len = 0;
94                e.stream().memset_zeros(&mut kvl.len_d)?;
95            }
96            for rl in sc.recur.iter_mut().flatten() {
97                e.stream().memset_zeros(&mut rl.conv_state)?;
98                e.stream().memset_zeros(&mut rl.ssm_state)?;
99                e.stream().memset_zeros(&mut rl.ssm_state_alt)?;
100            }
101            self.prime_chunk_captured(
102                e,
103                &x_in,
104                &pos_d,
105                bucket,
106                sc,
107                &len_d,
108                &mut lo_cell.borrow_mut(),
109                &mut hs_cell.borrow_mut(),
110            )
111        })?;
112        // Ends the &mut reborrows held by the capture cells before the buffers are
113        // moved into PrimeGraph below (the cells are RefCell<&mut _>, hence no Drop).
114        #[allow(clippy::drop_non_drop)] // allow: explicit end-of-borrow marker, see above
115        drop(scratch_cell);
116        #[allow(clippy::drop_non_drop)] // allow: explicit end-of-borrow marker, see above
117        drop(lo_cell);
118        #[allow(clippy::drop_non_drop)] // allow: explicit end-of-borrow marker, see above
119        drop(hs_cell);
120        // reclaim the private scratch (graph-baked) and restore the eager one
121        let private_scratch = e.f16_scratch_swap(prev_scratch);
122        let _ = CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED;
123        let _ = CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH;
124        Ok(PrimeGraph {
125            bucket,
126            graph,
127            _keeper: keeper,
128            private_scratch,
129            scratch,
130            x_in,
131            len_d,
132            logits_out,
133            h_seed_out,
134            n_embd,
135        })
136    }
137
138    /// Replay the graph for `tokens` (len <= bucket) and copy the outputs into `session`
139    /// (a FRESH cache: pos == 0). Returns host logits (the prefill_tick contract).
140    pub fn prime_graph_run(
141        &self,
142        e: &Engine,
143        pg: &mut PrimeGraph,
144        tokens: &[u32],
145        session: &mut Cache,
146    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
147        let t = tokens.len();
148        assert!(
149            t >= 2 && t <= pg.bucket,
150            "prime_graph_run: 2 <= T <= bucket"
151        );
152        assert!(session.pos == 0, "prime_graph_run: fresh sessions only");
153        let n_embd = pg.n_embd;
154        // graph inputs: embed rows + zeroed pad tail + true length (all OUTSIDE capture,
155        // so host-sourced writes are legal here)
156        let x = self.embed(e, tokens)?;
157        e.copy_into(&mut pg.x_in, 0, &x, t * n_embd)?;
158        if t < pg.bucket {
159            let mut tail = pg.x_in.slice_mut(t * n_embd..pg.bucket * n_embd);
160            e.stream().memset_zeros(&mut tail)?;
161        }
162        e.set_i32_one(&mut pg.len_d, t as i32)?;
163        pg.graph.launch()?;
164        // copy-out: quantized KV rows [0,T), conv rings, recurrent state
165        for (il, kvl) in pg.scratch.kv.iter().enumerate() {
166            let (Some(src), Some(dst)) = (kvl.as_ref(), session.kv[il].as_mut()) else {
167                continue;
168            };
169            let kb = t * src.k_tok_bytes;
170            let vb = t * src.v_tok_bytes;
171            e.stream()
172                .memcpy_dtod(&src.k.slice(0..kb), &mut dst.k.slice_mut(0..kb))?;
173            e.stream()
174                .memcpy_dtod(&src.v.slice(0..vb), &mut dst.v.slice_mut(0..vb))?;
175            dst.len = t;
176            e.set_i32_one(&mut dst.len_d, t as i32)?;
177        }
178        for (il, rl) in pg.scratch.recur.iter().enumerate() {
179            let (Some(src), Some(dst)) = (rl.as_ref(), session.recur[il].as_mut()) else {
180                continue;
181            };
182            let cn = src.conv_state.len();
183            let sn = src.ssm_state.len();
184            e.copy_into(&mut dst.conv_state, 0, &src.conv_state, cn)?;
185            e.copy_into(&mut dst.ssm_state, 0, &src.ssm_state, sn)?;
186        }
187        session.pos = t;
188        let logits = e.dtoh(&pg.logits_out)?;
189        let mut h_seed = e.uninit(n_embd)?;
190        let hn = pg.h_seed_out.len();
191        e.copy_into(&mut h_seed, 0, &pg.h_seed_out, hn)?;
192        Ok((logits, h_seed))
193    }
194}
memra_engine/prime_graph.rs

memra_engine/
prime_graph.rs
