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        use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
62        let n_embd = self.cfg.n_embd as usize;
63        let n_vocab = self.output.out_features();
64        let mut scratch = Cache::new(e, &self.cfg, bucket + 8)?;
65        let x_in = e.zeros(bucket * n_embd)?;
66        let pos_d = e.htod_i32(&(0..bucket as i32).collect::<Vec<_>>())?;
67        let len_d = e.htod_i32(&[bucket as i32])?;
68        let mut logits_out = e.uninit(n_vocab)?;
69        let mut h_seed_out = e.uninit(n_embd)?;
70
71        // capture with a PRIVATE f16 scratch resident (pre-sized to the trunk's largest
72        // GEMM input: m = bucket, in_f up to n_ff) so no eager call ever mutates the
73        // graph-baked buffers.
74        let n_ff_max = self
75            .layers
76            .iter()
77            .map(|l| match &l.ffn {
78                crate::hybrid::Ffn::Dense { ffn_gate, .. } => ffn_gate.out_features(),
79                _ => n_embd,
80            })
81            .max()
82            .unwrap_or(n_embd)
83            .max(n_embd);
84        let private = crate::f16_ffi::F16Scratch::with_capacity(e, bucket * n_ff_max * 2)?;
85        let prev_scratch = e.f16_scratch_swap(Some(private));
86        let scratch_cell = std::cell::RefCell::new(&mut scratch);
87        let lo_cell = std::cell::RefCell::new(&mut logits_out);
88        let hs_cell = std::cell::RefCell::new(&mut h_seed_out);
89        let (graph, keeper) = e.capture_graph_retained(|e| {
90            let sc: &mut Cache = &mut scratch_cell.borrow_mut();
91            for kvl in sc.kv.iter_mut().flatten() {
92                kvl.len = 0;
93                e.stream().memset_zeros(&mut kvl.len_d)?;
94            }
95            for rl in sc.recur.iter_mut().flatten() {
96                e.stream().memset_zeros(&mut rl.conv_state)?;
97                e.stream().memset_zeros(&mut rl.ssm_state)?;
98                e.stream().memset_zeros(&mut rl.ssm_state_alt)?;
99            }
100            self.prime_chunk_captured(
101                e,
102                &x_in,
103                &pos_d,
104                bucket,
105                sc,
106                &len_d,
107                &mut lo_cell.borrow_mut(),
108                &mut hs_cell.borrow_mut(),
109            )
110        })?;
111        drop(scratch_cell);
112        drop(lo_cell);
113        drop(hs_cell);
114        // reclaim the private scratch (graph-baked) and restore the eager one
115        let private_scratch = e.f16_scratch_swap(prev_scratch);
116        let _ = CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED;
117        let _ = CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH;
118        Ok(PrimeGraph {
119            bucket,
120            graph,
121            _keeper: keeper,
122            private_scratch,
123            scratch,
124            x_in,
125            len_d,
126            logits_out,
127            h_seed_out,
128            n_embd,
129        })
130    }
131
132    /// Replay the graph for `tokens` (len <= bucket) and copy the outputs into `session`
133    /// (a FRESH cache: pos == 0). Returns host logits (the prefill_tick contract).
134    pub fn prime_graph_run(
135        &self,
136        e: &Engine,
137        pg: &mut PrimeGraph,
138        tokens: &[u32],
139        session: &mut Cache,
140    ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn std::error::Error>> {
141        let t = tokens.len();
142        assert!(
143            t >= 2 && t <= pg.bucket,
144            "prime_graph_run: 2 <= T <= bucket"
145        );
146        assert!(session.pos == 0, "prime_graph_run: fresh sessions only");
147        let n_embd = pg.n_embd;
148        // graph inputs: embed rows + zeroed pad tail + true length (all OUTSIDE capture,
149        // so host-sourced writes are legal here)
150        let x = self.embed(e, tokens)?;
151        e.copy_into(&mut pg.x_in, 0, &x, t * n_embd)?;
152        if t < pg.bucket {
153            let mut tail = pg.x_in.slice_mut(t * n_embd..pg.bucket * n_embd);
154            e.stream().memset_zeros(&mut tail)?;
155        }
156        e.set_i32_one(&mut pg.len_d, t as i32)?;
157        pg.graph.launch()?;
158        // copy-out: quantized KV rows [0,T), conv rings, recurrent state
159        for (il, kvl) in pg.scratch.kv.iter().enumerate() {
160            let (Some(src), Some(dst)) = (kvl.as_ref(), session.kv[il].as_mut()) else {
161                continue;
162            };
163            let kb = t * src.k_tok_bytes;
164            let vb = t * src.v_tok_bytes;
165            e.stream()
166                .memcpy_dtod(&src.k.slice(0..kb), &mut dst.k.slice_mut(0..kb))?;
167            e.stream()
168                .memcpy_dtod(&src.v.slice(0..vb), &mut dst.v.slice_mut(0..vb))?;
169            dst.len = t;
170            e.set_i32_one(&mut dst.len_d, t as i32)?;
171        }
172        for (il, rl) in pg.scratch.recur.iter().enumerate() {
173            let (Some(src), Some(dst)) = (rl.as_ref(), session.recur[il].as_mut()) else {
174                continue;
175            };
176            let cn = src.conv_state.len();
177            let sn = src.ssm_state.len();
178            e.copy_into(&mut dst.conv_state, 0, &src.conv_state, cn)?;
179            e.copy_into(&mut dst.ssm_state, 0, &src.ssm_state, sn)?;
180        }
181        session.pos = t;
182        let logits = e.dtoh(&pg.logits_out)?;
183        let mut h_seed = e.uninit(n_embd)?;
184        let hn = pg.h_seed_out.len();
185        e.copy_into(&mut h_seed, 0, &pg.h_seed_out, hn)?;
186        Ok((logits, h_seed))
187    }
188}
memra_engine/prime_graph.rs

memra_engine/
prime_graph.rs
