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

memra_engine/
prime_graph.rs
