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