use crate::Engine;
use crate::hybrid::{HybridModel, Mixer};
use crate::hybrid_forward::WsSeg;
use cudarc::driver::{CudaGraph, CudaSlice};
use memra_kv::Cache;
use std::sync::atomic::Ordering;
type Res<T> = Result<T, Box<dyn std::error::Error>>;
enum RunPiece {
Graph {
graphs: [CudaGraph; 2],
#[allow(dead_code)]
segs: Vec<WsSeg>,
},
Middle(usize),
}
#[derive(Clone, Debug, PartialEq, Eq)]
enum PieceSpec {
Graph(Vec<WsSeg>),
Middle(usize),
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct RunSpec {
lo: usize,
hi: usize,
pieces: Vec<PieceSpec>,
}
struct RunGraph {
lo: usize,
hi: usize,
pieces: Vec<RunPiece>,
phase: usize,
checked: u32,
_keeper: Vec<Box<dyn std::any::Any + Send>>,
}
unsafe impl Send for RunGraph {}
unsafe impl Send for StageGraphs {}
struct StageGraphs {
dev: usize,
lo: usize,
hi: usize,
runs: Vec<RunGraph>,
x_io: CudaSlice<f32>,
x_out: CudaSlice<f32>,
ws: crate::hyper::HyperDecodeWs,
mixed: CudaSlice<f32>,
pos_d: CudaSlice<i32>,
spare: [CudaSlice<f32>; 2],
f16: Option<crate::f16_ffi::F16Scratch>,
next_pos: usize,
state_sig: Vec<LayerSig>,
launched_since_sync: bool,
}
struct StageBufs {
x_io: CudaSlice<f32>,
x_out: CudaSlice<f32>,
ws: crate::hyper::HyperDecodeWs,
mixed: CudaSlice<f32>,
pos_d: CudaSlice<i32>,
spare: [CudaSlice<f32>; 2],
f16: Option<crate::f16_ffi::F16Scratch>,
}
#[derive(Default)]
pub(crate) struct Glm5DecodeGraphs {
stages: Vec<StageGraphs>,
failed: Vec<(usize, usize, usize)>,
}
fn capturable_runs(
lo: usize,
hi: usize,
capturable: impl Fn(usize) -> bool,
) -> Vec<(usize, usize)> {
let mut runs = Vec::new();
let mut start: Option<usize> = None;
for il in lo..hi {
match (capturable(il), start) {
(true, None) => start = Some(il),
(false, Some(a)) => {
runs.push((a, il));
start = None;
}
_ => {}
}
}
if let Some(a) = start {
runs.push((a, hi));
}
runs
}
fn plan_runs(m: &HybridModel, lo: usize, hi: usize) -> Vec<RunSpec> {
let halves = crate::glm5_graph_mla_on() && Engine::mla_seg_ws_on();
if !halves {
return capturable_runs(
lo,
hi,
|il| matches!(&m.layers[il].mixer, Mixer::Kda(la) if la.tp.is_none()),
)
.into_iter()
.map(|(a, b)| RunSpec {
lo: a,
hi: b,
pieces: vec![PieceSpec::Graph((a..b).map(WsSeg::Layer).collect())],
})
.collect();
}
let mut runs: Vec<RunSpec> = Vec::new();
let mut start: Option<usize> = None;
let mut pieces: Vec<PieceSpec> = Vec::new();
let mut segs: Vec<WsSeg> = Vec::new();
fn close(
end: usize,
start: &mut Option<usize>,
pieces: &mut Vec<PieceSpec>,
segs: &mut Vec<WsSeg>,
runs: &mut Vec<RunSpec>,
) {
if !segs.is_empty() {
pieces.push(PieceSpec::Graph(std::mem::take(segs)));
}
if let Some(a) = start.take()
&& !pieces.is_empty()
{
runs.push(RunSpec {
lo: a,
hi: end,
pieces: std::mem::take(pieces),
});
}
pieces.clear();
}
for il in lo..hi {
let kda = matches!(&m.layers[il].mixer, Mixer::Kda(la) if la.tp.is_none());
let mla = halves && matches!(&m.layers[il].mixer, Mixer::Mla(mla) if mla.tp.is_none());
if kda {
start.get_or_insert(il);
segs.push(WsSeg::Layer(il));
} else if mla {
start.get_or_insert(il);
segs.push(WsSeg::MlaPre(il));
pieces.push(PieceSpec::Graph(std::mem::take(&mut segs)));
pieces.push(PieceSpec::Middle(il));
segs.push(WsSeg::MlaFfn(il));
} else {
close(il, &mut start, &mut pieces, &mut segs, &mut runs);
}
}
close(hi, &mut start, &mut pieces, &mut segs, &mut runs);
runs
}
fn kda_runs(m: &HybridModel, lo: usize, hi: usize) -> Vec<(usize, usize)> {
plan_runs(m, lo, hi).iter().map(|r| (r.lo, r.hi)).collect()
}
#[derive(Clone, Copy, PartialEq, Eq)]
struct LayerSig {
il: usize,
conv: u64,
p_lo: u64,
p_hi: u64,
}
fn recur_sig(e: &Engine, cache: &Cache, il: usize) -> Option<LayerSig> {
use cudarc::driver::DevicePtr;
let rl = cache.recur.get(il)?.as_ref()?;
let st = e.stream();
let (conv, _g0) = rl.conv_state.device_ptr(&st);
let (s, _g1) = rl.ssm_state.device_ptr(&st);
let (a, _g2) = rl.ssm_state_alt.device_ptr(&st);
Some(LayerSig {
il,
conv,
p_lo: s.min(a),
p_hi: s.max(a),
})
}
fn stage_sig(m: &HybridModel, e: &Engine, cache: &Cache, lo: usize, hi: usize) -> Vec<LayerSig> {
kda_runs(m, lo, hi)
.iter()
.flat_map(|(a, b)| *a..*b)
.filter_map(|il| recur_sig(e, cache, il))
.collect()
}
fn sig_diff(old: &[LayerSig], new: &[LayerSig]) -> Option<String> {
if old.len() != new.len() {
return Some(format!(
"layer-count {} -> {} (a captured layer lost or gained its recurrent slot)",
old.len(),
new.len()
));
}
for (o, n) in old.iter().zip(new.iter()) {
if o.il != n.il {
return Some(format!("layer order {} -> {}", o.il, n.il));
}
if o.conv != n.conv {
return Some(format!(
"layer {} conv_state 0x{:x} -> 0x{:x}",
o.il, o.conv, n.conv
));
}
if o.p_lo != n.p_lo || o.p_hi != n.p_hi {
return Some(format!(
"layer {} ssm pair {{0x{:x}, 0x{:x}}} -> {{0x{:x}, 0x{:x}}}",
o.il, o.p_lo, o.p_hi, n.p_lo, n.p_hi
));
}
}
None
}
fn selfcheck_n() -> u32 {
static N: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
*N.get_or_init(|| {
std::env::var("MEMRA_GLM5_GRAPH_SELFCHECK_N")
.ok()
.and_then(|v| v.trim().parse::<u32>().ok())
.filter(|&n| n >= 1)
.unwrap_or(1)
})
}
fn note(msg: &str) {
static SEEN: std::sync::OnceLock<std::sync::Mutex<std::collections::BTreeSet<String>>> =
std::sync::OnceLock::new();
let seen = SEEN.get_or_init(Default::default);
if seen.lock().unwrap().insert(msg.to_string()) {
eprintln!("[glm5-decode-graph] {msg}");
}
}
#[derive(Clone, Copy)]
struct CapCtx {
dev: usize,
lo: usize,
hi: usize,
run: usize,
runs: usize,
a: usize,
b: usize,
phase: usize,
recapture: bool,
}
impl std::fmt::Display for CapCtx {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"dev={} stage=[{}, {}) run={}/{} layers=[{}, {}) phase={} recapture={}",
self.dev,
self.lo,
self.hi,
self.run,
self.runs,
self.a,
self.b,
self.phase,
self.recapture
)
}
}
fn capture_status(e: &Engine) -> &'static str {
use cudarc::driver::sys;
let mut st = sys::CUstreamCaptureStatus::CU_STREAM_CAPTURE_STATUS_NONE;
let rc = unsafe { sys::cuStreamIsCapturing(e.stream().cu_stream(), &mut st) };
if rc != sys::CUresult::CUDA_SUCCESS {
return "query-failed";
}
match st {
sys::CUstreamCaptureStatus::CU_STREAM_CAPTURE_STATUS_NONE => "none",
sys::CUstreamCaptureStatus::CU_STREAM_CAPTURE_STATUS_ACTIVE => "ACTIVE",
sys::CUstreamCaptureStatus::CU_STREAM_CAPTURE_STATUS_INVALIDATED => "INVALIDATED",
}
}
fn free_mb(e: &Engine) -> String {
match e.ctx().mem_get_info() {
Ok((free, total)) => format!("{}/{}MB", free >> 20, total >> 20),
Err(_) => "?".to_string(),
}
}
fn capture_error(e: &Engine, ctx: &CapCtx, call: &str, err: &dyn std::fmt::Display) {
eprintln!(
"[glm5-decode-graph] capture-error: call={call} {ctx} stream_capture={} free={} \
ledger={} err={err}",
capture_status(e),
free_mb(e),
crate::glm5_sel_ledger::armed(),
);
}
fn step<T>(
e: &Engine,
ctx: &CapCtx,
call: &str,
r: Result<T, Box<dyn std::error::Error>>,
) -> Res<T> {
match r {
Ok(v) => Ok(v),
Err(err) => {
capture_error(e, ctx, call, &err);
Err(err)
}
}
}
fn capture_one<F>(e: &Engine, ctx: &CapCtx, mut body: F) -> Res<CudaGraph>
where
F: FnMut(&Engine) -> Res<()>,
{
use cudarc::driver::sys::{CUgraphInstantiate_flags, CUstreamCaptureMode};
let pre = capture_status(e);
if pre != "none" {
capture_error(
e,
ctx,
"pre-begin-status",
&format!("stream capture status is {pre}"),
);
return Err(format!("glm5 decode graph: stream capture status {pre} before begin").into());
}
step(
e,
ctx,
"synchronize(before begin_capture)",
e.stream().synchronize().map_err(Into::into),
)?;
let was_tracking = e.ctx().is_event_tracking();
if was_tracking {
unsafe { e.ctx().disable_event_tracking() };
}
let out = (|| -> Res<CudaGraph> {
step(
e,
ctx,
"cuStreamBeginCapture(RELAXED)",
e.stream()
.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)
.map_err(Into::into),
)?;
let body_res = body(e);
let ended = e
.stream()
.end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);
if let Err(err) = body_res {
capture_error(e, ctx, "capture body", &err);
return Err(err);
}
let graph = step(
e,
ctx,
"cuStreamEndCapture+cuGraphInstantiate",
ended.map_err(Into::into),
)?
.ok_or_else(|| {
capture_error(e, ctx, "cuStreamEndCapture", &"capture produced no graph");
"glm5 decode graph: capture produced no graph (stream was not capturing)"
})?;
step(e, ctx, "cuGraphUpload", graph.upload().map_err(Into::into))?;
Ok(graph)
})();
if was_tracking {
unsafe { e.ctx().enable_event_tracking() };
}
out
}
fn x_sum(e: &Engine, x: &CudaSlice<f32>) -> String {
match e.dtoh(x) {
Ok(v) => {
let mut h: u64 = 0xcbf2_9ce4_8422_2325;
let mut nz = 0usize;
let mut absmax = 0f32;
for f in &v {
h ^= f.to_bits() as u64;
h = h.wrapping_mul(0x100_0000_01b3);
if *f != 0.0 {
nz += 1;
}
if f.abs() > absmax {
absmax = f.abs();
}
}
format!("sum=0x{h:016x} nz={nz}/{} absmax={absmax:.6e}", v.len())
}
Err(err) => format!("sum=? ({err})"),
}
}
#[allow(clippy::too_many_arguments)] fn trace_seg(
e: &Engine,
dev: usize,
lo: usize,
hi: usize,
a: usize,
b: usize,
arm: &str,
pos: usize,
x: &CudaSlice<f32>,
) {
if !crate::glm5_graph_trace_on() {
return;
}
eprintln!(
"[glm5-graph-trace] pos={pos} dev={dev} stage=[{lo}, {hi}) seg=[{a}, {b}) arm={arm} {}",
x_sum(e, x)
);
}
impl HybridModel {
fn glm5_decode_graph_refusal(
&self,
e: &Engine,
cache: &Cache,
lo: usize,
hi: usize,
) -> Option<String> {
if self.hyper.is_none() {
return Some("model carries no HyperConnections topology".into());
}
if self.cfg.sigmoid_router().is_none() {
return Some(
"this trunk has no sigmoid router (the device-table MoE arm needs one)".into(),
);
}
if crate::glm5_graph_no_capture() {
return Some(
"MEMRA_GLM5_GRAPH_NO_CAPTURE: capture disabled while the door's device-table MoE \
arm stays engaged (the half of the bisect MEMRA_GLM5_GRAPH_HOST_MOE could not \
supply, since that one turns off both enablers at once)"
.into(),
);
}
if crate::glm5_graph_host_moe() {
return Some(
"MEMRA_GLM5_GRAPH_HOST_MOE forces the host-oracle MoE: its per-layer readback and \
stream drain cannot live inside a capture region (bisect arm)"
.into(),
);
}
if !crate::htod_diet_on() {
return Some(
"MEMRA_HTOD_DIET is off: the shared expert still uploads a pageable constant \
per MoE layer, which a capture region refuses"
.into(),
);
}
if !crate::hybrid_forward::sigmoid_router_enabled() {
return Some(
"MEMRA_SIG_ROUTER=0 selects the host oracle (no device selection to capture)"
.into(),
);
}
if crate::moesd::capture_active() || memra_reference::hidden_trace::enabled() {
return Some("a host-visible route/hidden observer is armed".into());
}
for env in [
"MEMRA_MOE_STATS",
"MEMRA_MOE_TRACE",
"MEMRA_MOE_WEIGHT_TRACE",
"MEMRA_MOE_INPUT_TRACE_DIR",
"MEMRA_SIG_ROUTER_LOGIT_TRACE",
"MEMRA_MOE_SEL_DUMP",
] {
if std::env::var_os(env).is_some() {
return Some(format!(
"{env} is armed (its consumer reads the selection on the host)"
));
}
}
if e.ctx().is_event_tracking() {
return Some(
"cudarc event tracking is on (MEMRA_EVT); capture refuses cross-stream waits"
.into(),
);
}
for il in lo..hi {
if let Mixer::Kda(la) = &self.layers[il].mixer
&& la.tp.is_some()
{
return Some(format!("layer {il} is glm5-TP sharded"));
}
if cache.recur.get(il).is_none() {
return Some(format!(
"layer {il} has no recurrent-state slot in this cache"
));
}
}
let runs = kda_runs(self, lo, hi);
if runs.is_empty() {
return Some(format!("stage [{lo}, {hi}) holds no KDA layer to capture"));
}
for (a, b) in runs {
for il in a..b {
if !HybridModel::glm5_t1_dev_moe_ready(e, &self.layers[il], &self.cfg, il) {
return Some(format!(
"layer {il}: the T=1 device-table MoE arm would not fire (it needs \
slab-local uniform q8 experts, a PRE clamp, and n_used <= 8)"
));
}
}
}
None
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn hyper_range_decode_eager_traced(
&self,
e: &Engine,
topology: &crate::hyper::HyperTopology,
mut x: CudaSlice<f32>,
lo: usize,
hi: usize,
pos_d: &CudaSlice<i32>,
pos: usize,
cache: &mut Cache,
) -> Res<CudaSlice<f32>> {
let dev = e.ctx().ordinal();
let runs = kda_runs(self, lo, hi);
let mut cursor = lo;
for (a, b) in runs {
if a > cursor {
x = self.hyper_range_decode_eager(e, topology, x, cursor, a, pos_d, pos, cache)?;
trace_seg(e, dev, lo, hi, cursor, a, "eager-gap", pos, &x);
}
x = self.hyper_range_decode_eager(e, topology, x, a, b, pos_d, pos, cache)?;
trace_seg(e, dev, lo, hi, a, b, "eager-run", pos, &x);
cursor = b;
}
if cursor < hi {
x = self.hyper_range_decode_eager(e, topology, x, cursor, hi, pos_d, pos, cache)?;
trace_seg(e, dev, lo, hi, cursor, hi, "eager-gap", pos, &x);
}
Ok(x)
}
pub(crate) fn glm5_decode_graph_ready(
&self,
e: &Engine,
cache: &Cache,
lo: usize,
hi: usize,
) -> bool {
if cache
.glm5_decode_graph
.as_ref()
.and_then(|b| b.downcast_ref::<Glm5DecodeGraphs>())
.is_some_and(|p| {
p.failed
.iter()
.any(|&(d, l, h)| d == e.ctx().ordinal() && l == lo && h == hi)
})
{
return false;
}
if let Some(why) = self.glm5_decode_graph_refusal(e, cache, lo, hi) {
note(&format!("eager: {why}"));
return false;
}
if !crate::spec::graph_launch_headroom_ok(e) {
static SUSPENDED: std::sync::Once = std::sync::Once::new();
SUSPENDED.call_once(|| crate::spec::graph_replay_suspended_note("glm5-decode-graph"));
return false;
}
true
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn hyper_range_decode_graphed(
&self,
e: &Engine,
topology: &crate::hyper::HyperTopology,
mut x: CudaSlice<f32>,
lo: usize,
hi: usize,
pos_d: &CudaSlice<i32>,
pos: usize,
cache: &mut Cache,
) -> Res<CudaSlice<f32>> {
let dev = e.ctx().ordinal();
let n_embd = self.cfg.n_embd as usize;
let width = topology.streams * n_embd;
let live_sig = stage_sig(self, e, cache, lo, hi);
let mut reuse: Option<StageBufs> = None;
let mut latch_eager = false;
let recapture_armed = crate::glm5_graph_recapture_on();
let mut need_capture = {
let pool = self.glm5_graph_pool(cache);
match pool
.stages
.iter()
.position(|s| s.dev == dev && s.lo == lo && s.hi == hi)
{
Some(i) => {
let st = &pool.stages[i];
let stale_pos = st.next_pos != pos;
let diff = sig_diff(&st.state_sig, &live_sig);
if stale_pos || diff.is_some() {
eprintln!(
"[glm5-decode-graph] {} dev={dev} stage=[{lo}, {hi}) pos={pos} \
expected_pos={} stale_pos={stale_pos} launched_since_sync={} \
stream_capture={} free={} sig_diff={}",
if recapture_armed {
"re-capture"
} else {
"eager-latch (MEMRA_GLM5_GRAPH_RECAPTURE is off; this stage runs \
eager for the rest of the session and the walk stays \
byte-identical)"
},
st.next_pos,
st.launched_since_sync,
capture_status(e),
free_mb(e),
diff.as_deref().unwrap_or("none"),
);
if recapture_armed {
true
} else {
pool.failed.push((dev, lo, hi));
latch_eager = true;
false
}
} else {
false
}
}
None => true,
}
};
if recapture_armed && !latch_eager && need_capture {
let present = {
let pool = self.glm5_graph_pool(cache);
pool.stages
.iter()
.position(|s| s.dev == dev && s.lo == lo && s.hi == hi)
};
if let Some(i) = present {
if let Err(err) = e.stream().synchronize() {
note(&format!(
"eager from here on stage=[{lo}, {hi}) dev={dev}: the pre-teardown drain \
failed ({err}), so the stale execs are left alone"
));
self.glm5_graph_pool(cache).failed.push((dev, lo, hi));
return self
.hyper_range_decode_eager(e, topology, x, lo, hi, pos_d, pos, cache);
}
let stale = self.glm5_graph_pool(cache).stages.remove(i);
crate::GLM5_DECODE_GRAPH_RECAPTURES.fetch_add(1, Ordering::Relaxed);
let StageGraphs {
runs,
x_io,
x_out,
ws,
mixed,
pos_d,
spare,
f16,
..
} = stale;
reuse = Some(StageBufs {
x_io,
x_out,
ws,
mixed,
pos_d,
spare,
f16,
});
drop(runs);
need_capture = true;
}
}
if latch_eager {
return self.hyper_range_decode_eager(e, topology, x, lo, hi, pos_d, pos, cache);
}
if need_capture {
if let Err(err) = self
.glm5_warm_runs_before_capture(e, topology, &x, lo, hi, width, pos_d, pos, cache)
{
note(&format!(
"eager from here on stage=[{lo}, {hi}) dev={dev}: the pre-capture warm walk failed ({err})"
));
self.glm5_graph_pool(cache).failed.push((dev, lo, hi));
return self.hyper_range_decode_eager(e, topology, x, lo, hi, pos_d, pos, cache);
}
}
if need_capture
&& let Err(err) =
self.glm5_capture_stage(e, topology, lo, hi, cache, dev, width, pos, reuse)
{
note(&format!(
"eager from here on stage=[{lo}, {hi}) dev={dev}: capture failed ({err})"
));
let _ = e.stream().synchronize();
self.glm5_graph_pool(cache).failed.push((dev, lo, hi));
return self.hyper_range_decode_eager(e, topology, x, lo, hi, pos_d, pos, cache);
}
{
let pool = self.glm5_graph_pool(cache);
if let Some(st) = pool
.stages
.iter_mut()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
{
e.i32_mirror_store(&mut st.pos_d, pos as i32)?;
}
}
let runs: Vec<(usize, usize)> = kda_runs(self, lo, hi);
let mut cursor = lo;
for (a, b) in runs {
if a > cursor {
x = self.hyper_range_decode_eager(e, topology, x, cursor, a, pos_d, pos, cache)?;
trace_seg(e, dev, lo, hi, cursor, a, "graph-gap", pos, &x);
}
let unchecked = self
.glm5_graph_pool(cache)
.stages
.iter()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
.and_then(|st| st.runs.iter().find(|r| r.lo == a && r.hi == b))
.is_some_and(|r| r.checked < selfcheck_n());
if unchecked {
let (x_out, ok) = self.glm5_selfcheck_run(
e, topology, dev, lo, hi, a, b, x, width, pos_d, pos, cache,
)?;
x = x_out;
if !ok {
x = self.hyper_range_decode_eager(e, topology, x, b, hi, pos_d, pos, cache)?;
return Ok(x);
}
} else {
x = self.glm5_replay_run(e, dev, lo, hi, a, b, x, width, pos_d, cache)?;
}
trace_seg(e, dev, lo, hi, a, b, "graph-run", pos, &x);
cursor = b;
}
if cursor < hi {
x = self.hyper_range_decode_eager(e, topology, x, cursor, hi, pos_d, pos, cache)?;
trace_seg(e, dev, lo, hi, cursor, hi, "graph-gap", pos, &x);
}
{
let pool = self.glm5_graph_pool(cache);
if let Some(st) = pool
.stages
.iter_mut()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
{
st.next_pos = pos + 1;
}
}
Ok(x)
}
fn glm5_graph_pool<'a>(&self, cache: &'a mut Cache) -> &'a mut Glm5DecodeGraphs {
if cache
.glm5_decode_graph
.as_ref()
.and_then(|b| b.downcast_ref::<Glm5DecodeGraphs>())
.is_none()
{
cache.glm5_decode_graph = Some(Box::new(Glm5DecodeGraphs::default()));
}
cache
.glm5_decode_graph
.as_mut()
.and_then(|b| b.downcast_mut::<Glm5DecodeGraphs>())
.expect("just installed")
}
#[allow(clippy::too_many_arguments)] fn glm5_capture_stage(
&self,
e: &Engine,
topology: &crate::hyper::HyperTopology,
lo: usize,
hi: usize,
cache: &mut Cache,
dev: usize,
width: usize,
pos: usize,
reuse: Option<StageBufs>,
) -> Res<()> {
let n_embd = self.cfg.n_embd as usize;
let runs = plan_runs(self, lo, hi);
let recapture = reuse.is_some();
let alloc_ctx = CapCtx {
dev,
lo,
hi,
run: 0,
runs: runs.len(),
a: lo,
b: hi,
phase: 0,
recapture,
};
let (x_io, x_out, ws, mixed, pos_d_buf, spare, f16) = match reuse {
Some(b) => (b.x_io, b.x_out, b.ws, b.mixed, b.pos_d, b.spare, b.f16),
None => (
step(e, &alloc_ctx, "alloc(x_io)", e.zeros(width))?,
step(e, &alloc_ctx, "alloc(x_out)", e.zeros(width))?,
step(
e,
&alloc_ctx,
"alloc(HyperDecodeWs)",
crate::hyper::HyperDecodeWs::new(e, topology, n_embd),
)?,
step(e, &alloc_ctx, "alloc(mixed)", e.zeros(n_embd))?,
step(
e,
&alloc_ctx,
"htod_i32(pos_d)",
e.htod_i32(&[pos as i32]).map(Some),
)?
.expect("htod returned a buffer"),
[
step(e, &alloc_ctx, "alloc(spare)", e.zeros(1))?,
step(e, &alloc_ctx, "alloc(spare)", e.zeros(1))?,
],
None,
),
};
let f16 = match f16 {
Some(f) => f,
None => step(
e,
&alloc_ctx,
"alloc(F16Scratch)",
crate::f16_ffi::F16Scratch::with_capacity(e, (4 << 20).max(n_embd * 16)),
)?,
};
let mut stage = StageGraphs {
dev,
lo,
hi,
runs: Vec::with_capacity(runs.len()),
x_io,
x_out,
ws,
mixed,
pos_d: pos_d_buf,
spare,
f16: None,
next_pos: pos,
state_sig: runs
.iter()
.flat_map(|r| r.lo..r.hi)
.filter_map(|il| recur_sig(e, cache, il))
.collect(),
launched_since_sync: false,
};
if crate::glm5_sel_ledger::armed()
&& let Some(moe) = self.cfg.moe.as_ref()
{
for r in &runs {
for il in r.lo..r.hi {
crate::glm5_sel_ledger::prearm(e, il as u16, moe.expert_used_count as usize)?;
}
}
}
let prev_f16 = e.f16_scratch_swap(Some(f16));
let captured = (|| -> Res<()> {
crate::GLM5_GRAPH_CAPTURE_OPEN.store(true, Ordering::Relaxed);
let mut ctx = CapCtx {
dev,
lo,
hi,
run: 0,
runs: runs.len(),
a: 0,
b: 0,
phase: 0,
recapture,
};
step(
e,
&ctx,
"i32_mirror_store(pos_d)",
e.i32_mirror_store(&mut stage.pos_d, pos as i32),
)?;
for (ri, run) in runs.iter().enumerate() {
let (a, b) = (run.lo, run.hi);
ctx.run = ri;
ctx.a = a;
ctx.b = b;
let mut per_piece: Vec<Vec<CudaGraph>> = Vec::new();
for phase in 0..2 {
ctx.phase = phase;
let mut gi = 0usize;
for piece in &run.pieces {
let PieceSpec::Graph(segs) = piece else {
continue;
};
let x_cell = std::cell::RefCell::new(&mut stage.x_io);
let out_cell = std::cell::RefCell::new(&mut stage.x_out);
let ws_cell = std::cell::RefCell::new(&mut stage.ws);
let cache_cell = std::cell::RefCell::new(&mut *cache);
let mixed_ref: &CudaSlice<f32> = &stage.mixed;
let pos_d_ref: &CudaSlice<i32> = &stage.pos_d;
let g = capture_one(e, &ctx, |e| {
self.hyper_range_decode_ws_segments(
e,
topology,
&mut x_cell.borrow_mut(),
segs,
pos_d_ref,
pos,
&mut cache_cell.borrow_mut(),
&mut ws_cell.borrow_mut(),
Some(mixed_ref),
)?;
let live = x_cell.borrow();
e.copy_into(&mut out_cell.borrow_mut(), 0, &live, width)
})?;
if phase == 0 {
per_piece.push(vec![g]);
} else {
per_piece[gi].push(g);
}
gi += 1;
}
}
let mut it = per_piece.into_iter();
let mut pieces: Vec<RunPiece> = Vec::with_capacity(run.pieces.len());
for piece in &run.pieces {
match piece {
PieceSpec::Graph(segs) => {
let mut v = it.next().expect("captured both phases");
let g0 = v.remove(0);
let g1 = v.remove(0);
pieces.push(RunPiece::Graph {
graphs: [g0, g1],
segs: segs.clone(),
});
}
PieceSpec::Middle(il) => pieces.push(RunPiece::Middle(*il)),
}
}
let keeper = std::mem::take(&mut *e.glm5_graph_keep().lock().unwrap());
stage.runs.push(RunGraph {
lo: a,
hi: b,
pieces,
phase: 0,
checked: 0,
_keeper: keeper,
});
crate::GLM5_DECODE_GRAPH_LAYERS.fetch_add((b - a) as u64, Ordering::Relaxed);
}
Ok(())
})();
crate::GLM5_GRAPH_CAPTURE_OPEN.store(false, Ordering::Relaxed);
stage.f16 = e.f16_scratch_swap(prev_f16);
captured?;
crate::GLM5_DECODE_GRAPH_CAPTURES.fetch_add(1, Ordering::Relaxed);
let halves: usize = stage
.runs
.iter()
.map(|r| {
r.pieces
.iter()
.filter(|p| matches!(p, RunPiece::Middle(_)))
.count()
})
.sum();
crate::GLM5_DECODE_GRAPH_MLA_HALVES.fetch_add(halves as u64, Ordering::Relaxed);
let line = format!(
"engaged dev={dev} stage=[{lo}, {hi}) runs={} captured_layers={} mla_halves={halves} \
recapture={recapture} free={} (2 ping-pong phases each; {})",
stage.runs.len(),
stage.runs.iter().map(|r| r.hi - r.lo).sum::<usize>(),
free_mb(e),
if halves > 0 {
"MLA layers captured in halves around an eager middle, MEMRA_GLM5_GRAPH_MLA=1"
} else {
"MLA/DSA layers stay eager"
},
);
if recapture {
eprintln!("[glm5-decode-graph] {line}");
} else {
note(&line);
}
self.glm5_graph_pool(cache).stages.push(stage);
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn glm5_warm_runs_before_capture(
&self,
e: &Engine,
topology: &crate::hyper::HyperTopology,
x: &CudaSlice<f32>,
lo: usize,
hi: usize,
width: usize,
pos_d: &CudaSlice<i32>,
pos: usize,
cache: &mut Cache,
) -> Res<()> {
use cudarc::driver::DevicePtr;
let runs = kda_runs(self, lo, hi);
for (a, b) in runs {
struct Snap {
il: usize,
conv: CudaSlice<f32>,
ssm: CudaSlice<f32>,
alt: CudaSlice<f32>,
ssm_ptr: u64,
}
let mla_lens: Vec<(usize, usize)> = (a..b)
.filter_map(|il| cache.latent[il].as_ref().map(|l| (il, l.len)))
.collect();
let mut snaps: Vec<Snap> = Vec::with_capacity(b - a);
for il in a..b {
let Some(rl) = cache.recur[il].as_ref() else {
continue;
};
let ssm_ptr = {
let st = e.stream();
let (p, _g) = rl.ssm_state.device_ptr(&st);
p
};
let mut conv = e.uninit(rl.conv_state.len())?;
let mut ssm = e.uninit(rl.ssm_state.len())?;
let mut alt = e.uninit(rl.ssm_state_alt.len())?;
e.copy_into(&mut conv, 0, &rl.conv_state, rl.conv_state.len())?;
e.copy_into(&mut ssm, 0, &rl.ssm_state, rl.ssm_state.len())?;
e.copy_into(&mut alt, 0, &rl.ssm_state_alt, rl.ssm_state_alt.len())?;
snaps.push(Snap {
il,
conv,
ssm,
alt,
ssm_ptr,
});
}
let mut xc = e.uninit(width)?;
e.copy_into(&mut xc, 0, x, width)?;
let _ = self.hyper_range_decode_eager(e, topology, xc, a, b, pos_d, pos, cache)?;
for sn in &snaps {
let rl = cache.recur[sn.il].as_mut().expect("snapshotted above");
let p = {
let st = e.stream();
let (p, _g) = rl.ssm_state.device_ptr(&st);
p
};
if p != sn.ssm_ptr {
std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
}
e.copy_into(&mut rl.conv_state, 0, &sn.conv, sn.conv.len())?;
e.copy_into(&mut rl.ssm_state, 0, &sn.ssm, sn.ssm.len())?;
e.copy_into(&mut rl.ssm_state_alt, 0, &sn.alt, sn.alt.len())?;
}
for (il, len) in &mla_lens {
let l = cache.latent[*il].as_mut().expect("snapshotted above");
l.len = *len;
e.i32_mirror_store(&mut l.len_d, *len as i32)?;
}
note(&format!(
"warmed run [{a}, {b}) before capture (one eager pass on a copy of the input, \
state restored): lazily built per-weight caches now exist"
));
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn glm5_selfcheck_run(
&self,
e: &Engine,
topology: &crate::hyper::HyperTopology,
dev: usize,
lo: usize,
hi: usize,
a: usize,
b: usize,
x: CudaSlice<f32>,
width: usize,
pos_d: &CudaSlice<i32>,
pos: usize,
cache: &mut Cache,
) -> Res<(CudaSlice<f32>, bool)> {
struct Snap {
il: usize,
conv: CudaSlice<f32>,
ssm: CudaSlice<f32>,
alt: CudaSlice<f32>,
ssm_ptr: u64,
}
fn ptr_of(e: &Engine, s: &CudaSlice<f32>) -> u64 {
use cudarc::driver::DevicePtr;
let st = e.stream();
let (p, _g) = s.device_ptr(&st);
p
}
let mut snaps: Vec<Snap> = Vec::with_capacity(b - a);
let mla_lens: Vec<(usize, usize)> = (a..b)
.filter_map(|il| cache.latent[il].as_ref().map(|l| (il, l.len)))
.collect();
for il in a..b {
let Some(rl) = cache.recur[il].as_ref() else {
continue;
};
let ssm_ptr = ptr_of(e, &rl.ssm_state);
let mut conv = e.uninit(rl.conv_state.len())?;
let mut ssm = e.uninit(rl.ssm_state.len())?;
let mut alt = e.uninit(rl.ssm_state_alt.len())?;
e.copy_into(&mut conv, 0, &rl.conv_state, rl.conv_state.len())?;
e.copy_into(&mut ssm, 0, &rl.ssm_state, rl.ssm_state.len())?;
e.copy_into(&mut alt, 0, &rl.ssm_state_alt, rl.ssm_state_alt.len())?;
snaps.push(Snap {
il,
conv,
ssm,
alt,
ssm_ptr,
});
}
let restore = |e: &Engine, cache: &mut Cache, snaps: &[Snap]| -> Res<()> {
for sn in snaps {
let rl = cache.recur[sn.il].as_mut().expect("snapshotted above");
if ptr_of(e, &rl.ssm_state) != sn.ssm_ptr {
std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
}
e.copy_into(&mut rl.conv_state, 0, &sn.conv, sn.conv.len())?;
e.copy_into(&mut rl.ssm_state, 0, &sn.ssm, sn.ssm.len())?;
e.copy_into(&mut rl.ssm_state_alt, 0, &sn.alt, sn.alt.len())?;
}
for (il, len) in &mla_lens {
let l = cache.latent[*il].as_mut().expect("snapshotted above");
l.len = *len;
e.i32_mirror_store(&mut l.len_d, *len as i32)?;
}
Ok(())
};
let nonfinite = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
let mut nf_state = 0usize;
let mut nf_state_first = String::new();
for sn in &snaps {
for (name, buf) in [
("conv_state", &sn.conv),
("ssm_state", &sn.ssm),
("ssm_state_alt", &sn.alt),
] {
let n = nonfinite(&e.dtoh(buf)?);
if n > 0 && nf_state_first.is_empty() {
nf_state_first = format!("layer {} {name} {n}/{}", sn.il, buf.len());
}
nf_state += n;
}
}
let (nf_pool, nf_pool_first) = {
let pool = self.glm5_graph_pool(cache);
let mut total = 0usize;
let mut first = String::new();
if let Some(st) = pool
.stages
.iter()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
{
let bufs: [(&str, &CudaSlice<f32>); 10] = [
("x_io", &st.x_io),
("x_out", &st.x_out),
("ws.h", &st.ws.h),
("ws.y", &st.ws.y),
("ws.z", &st.ws.z),
("ws.xb", &st.ws.xb),
("ws.mixes", &st.ws.mixes),
("ws.pre", &st.ws.pre),
("ws.post", &st.ws.post),
("ws.comb", &st.ws.comb),
];
for (name, buf) in bufs {
let n = nonfinite(&e.dtoh(buf)?);
if n > 0 && first.is_empty() {
first = format!("{name} {n}/{}", buf.len());
}
total += n;
}
}
(total, first)
};
let mut x_keep = e.uninit(width)?;
e.copy_into(&mut x_keep, 0, &x, width)?;
let mut x_copy = e.uninit(width)?;
e.copy_into(&mut x_copy, 0, &x, width)?;
let x_ref = self.hyper_range_decode_eager(e, topology, x_copy, a, b, pos_d, pos, cache)?;
let h_ref = e.dtoh(&x_ref)?;
struct Post {
il: usize,
conv: Vec<f32>,
ssm: Vec<f32>,
alt: Vec<f32>,
ssm_ptr: u64,
}
let mut post: Vec<Post> = Vec::with_capacity(b - a);
for il in a..b {
let Some(rl) = cache.recur[il].as_ref() else {
continue;
};
post.push(Post {
il,
conv: e.dtoh(&rl.conv_state)?,
ssm: e.dtoh(&rl.ssm_state)?,
alt: e.dtoh(&rl.ssm_state_alt)?,
ssm_ptr: ptr_of(e, &rl.ssm_state),
});
}
restore(e, cache, &snaps)?;
let (nf_ews, nf_ews_first) = match e.hyper_ws_take() {
Some(ws) => {
let mut total = 0usize;
let mut first = String::new();
for (name, buf) in [
("ws.h", &ws.h),
("ws.y", &ws.y),
("ws.z", &ws.z),
("ws.xb", &ws.xb),
("ws.mixes", &ws.mixes),
("ws.pre", &ws.pre),
("ws.post", &ws.post),
("ws.comb", &ws.comb),
] {
let n = nonfinite(&e.dtoh(buf)?);
if n > 0 && first.is_empty() {
first = format!("{name} {n}/{}", buf.len());
}
total += n;
}
e.hyper_ws_put(ws);
(total, first)
}
None => (0, "no pooled ws".to_string()),
};
let mut x_copy2 = e.uninit(width)?;
e.copy_into(&mut x_copy2, 0, &x_keep, width)?;
crate::hybrid_forward::HC_WS_FORCE_PLAIN.store(true, std::sync::atomic::Ordering::Relaxed);
let plain = self.hyper_range_decode_eager(e, topology, x_copy2, a, b, pos_d, pos, cache);
crate::hybrid_forward::HC_WS_FORCE_PLAIN.store(false, std::sync::atomic::Ordering::Relaxed);
let x_plain = plain?;
let h_plain = e.dtoh(&x_plain)?;
restore(e, cache, &snaps)?;
let nf_plain = nonfinite(&h_plain);
let plain_vs_ws = match h_ref
.iter()
.zip(h_plain.iter())
.position(|(r, p)| r.to_bits() != p.to_bits())
{
None => "bit-identical to the workspace eager".to_string(),
Some(i) => format!(
"differs from the workspace eager at element {i} (ws={:e} plain={:e})",
h_ref[i], h_plain[i]
),
};
let (k, phase) = {
let pool = self.glm5_graph_pool(cache);
pool.stages
.iter()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
.and_then(|st| st.runs.iter().find(|r| r.lo == a && r.hi == b))
.map(|r| (r.checked + 1, r.phase))
.unwrap_or((1, 0))
};
let x_rep = self.glm5_replay_run(e, dev, lo, hi, a, b, x, width, pos_d, cache)?;
let h_rep = e.dtoh(&x_rep)?;
let nonfinite = |v: &[f32]| v.iter().filter(|x| !x.is_finite()).count();
let (nf_in, nf_ref, nf_rep) = {
let h_in = e.dtoh(&x_keep)?;
(nonfinite(&h_in), nonfinite(&h_ref), nonfinite(&h_rep))
};
let layer_bisect = if nf_in == 0 && nf_ref > 0 {
let mut report = String::new();
let mut xl = e.uninit(width)?;
e.copy_into(&mut xl, 0, &x_keep, width)?;
let mut first_bad: Option<(usize, CudaSlice<f32>)> = None;
for il in a..b {
let mut xin = e.uninit(width)?;
e.copy_into(&mut xin, 0, &xl, width)?;
let out =
self.hyper_range_decode_eager(e, topology, xl, il, il + 1, pos_d, pos, cache)?;
let n = nonfinite(&e.dtoh(&out)?);
report.push_str(&format!(" L{il}:{n}"));
if n > 0 {
first_bad = Some((il, xin));
break;
}
xl = out;
}
restore(e, cache, &snaps)?;
let sub = match first_bad {
Some((il, xin)) => {
let n_embd = self.cfg.n_embd as usize;
let eps = self.cfg.rms_eps;
let layer = &self.layers[il];
let mut steps = String::new();
if let Some(hyper) = layer.hyper.as_ref() {
let (y, mix) =
crate::hyper::pre(e, topology, &hyper.attn, &xin, 1, n_embd)?;
steps.push_str(&format!(" pre_attn.y:{}", nonfinite(&e.dtoh(&y)?)));
let mut h = e.uninit(n_embd)?;
e.rms_norm(&y, layer.attn_norm.float_data(), &mut h, n_embd, 1, eps)?;
steps.push_str(&format!(" attn_norm.h:{}", nonfinite(&e.dtoh(&h)?)));
let mixed = match &layer.mixer {
crate::hybrid::Mixer::Kda(la) if la.tp.is_none() => {
Some(crate::kda::kda_decode_cached(e, la, &h, eps, cache, il)?)
}
_ => None,
};
if let Some(mixed) = mixed.as_ref() {
steps.push_str(&format!(" kda_mixer:{}", nonfinite(&e.dtoh(mixed)?)));
let x1 = crate::hyper::post(e, topology, mixed, &xin, &mix, 1, n_embd)?;
steps.push_str(&format!(" post_attn.x:{}", nonfinite(&e.dtoh(&x1)?)));
let (y2, mix2) =
crate::hyper::pre(e, topology, &hyper.mlp, &x1, 1, n_embd)?;
steps.push_str(&format!(" pre_mlp.y:{}", nonfinite(&e.dtoh(&y2)?)));
let mut z = e.uninit(n_embd)?;
e.rms_norm(
&y2,
layer.post_attn_norm.float_data(),
&mut z,
n_embd,
1,
eps,
)?;
steps.push_str(&format!(" mlp_norm.z:{}", nonfinite(&e.dtoh(&z)?)));
let ffn_out =
self.hyper_ffn_branch(e, layer, &z, 1, il, false, None)?;
steps.push_str(&format!(" ffn_out:{}", nonfinite(&e.dtoh(&ffn_out)?)));
let x2 =
crate::hyper::post(e, topology, &ffn_out, &x1, &mix2, 1, n_embd)?;
steps.push_str(&format!(" post_mlp.x:{}", nonfinite(&e.dtoh(&x2)?)));
} else {
steps.push_str(" (mixer is not a plain KDA layer; sub-steps skipped)");
}
}
restore(e, cache, &snaps)?;
format!("; inside L{il} sub-steps non-finite:{steps}")
}
None => String::new(),
};
format!("; per-layer non-finite after each layer:{report}{sub}")
} else {
String::new()
};
if nf_in + nf_ref + nf_rep > 0 {
eprintln!(
"[glm5-decode-graph] SELF-CHECK FAILED dev={dev} stage=[{lo}, {hi}) \
run=[{a}, {b}) pos={pos} replay={k} phase={phase}: NON-FINITE data: input \
{nf_in}/{width}, eager output {nf_ref}/{width}, replay output {nf_rep}/{width}; \
BEFORE the walk: recurrent state non-finite {nf_state} (first: {}), stage pool \
non-finite {nf_pool} (first: {}), engine ws non-finite {nf_ews} (first: {}); PLAIN \
eager twin: non-finite {nf_plain}/{width}, {plain_vs_ws}{layer_bisect} \
(a NaN-identical compare proves nothing); the stage is latched EAGER for this \
session and this token is recomputed eager. The door did NOT engage.{}",
if nf_state_first.is_empty() {
"none"
} else {
&nf_state_first
},
if nf_pool_first.is_empty() {
"none"
} else {
&nf_pool_first
},
if nf_ews_first.is_empty() {
"none"
} else {
&nf_ews_first
},
if nf_in > 0 {
" The INPUT was already poisoned: the defect is upstream of this stage."
} else if nf_state > 0 {
" The recurrent STATE was already poisoned before this run's walk."
} else if nf_pool > 0 {
" The stage POOL (workspace / x cells) was already poisoned before this run's walk."
} else {
" Input, state and pool were finite: the eager walk itself produced the poison."
}
);
restore(e, cache, &snaps)?;
let x_real =
self.hyper_range_decode_eager(e, topology, x_keep, a, b, pos_d, pos, cache)?;
self.glm5_graph_pool(cache).failed.push((dev, lo, hi));
return Ok((x_real, false));
}
let first_diff = h_ref
.iter()
.zip(h_rep.iter())
.position(|(r, p)| r.to_bits() != p.to_bits())
.map(|i| {
format!(
"OUTPUT element {i}/{width} (eager={:e} replay={:e})",
h_ref[i], h_rep[i]
)
});
let state_diff = if first_diff.is_some() {
None
} else {
let mut found = None;
for p in &post {
let rl = cache.recur[p.il].as_ref().expect("snapshotted above");
if ptr_of(e, &rl.ssm_state) != p.ssm_ptr {
found = Some(format!(
"layer {} ssm ROLE: the replay left `ssm_state` at a different parity \
than the eager walk",
p.il
));
break;
}
let bufs: [(&str, Vec<f32>, &Vec<f32>); 3] = [
("conv_state", e.dtoh(&rl.conv_state)?, &p.conv),
("ssm_state", e.dtoh(&rl.ssm_state)?, &p.ssm),
("ssm_state_alt", e.dtoh(&rl.ssm_state_alt)?, &p.alt),
];
for (name, got, want) in &bufs {
if let Some(i) = got
.iter()
.zip(want.iter())
.position(|(g, w)| g.to_bits() != w.to_bits())
{
found = Some(format!(
"layer {} {name} element {i}/{} (eager={:e} replay={:e})",
p.il,
want.len(),
want[i],
got[i]
));
break;
}
}
if found.is_some() {
break;
}
}
found
};
match first_diff.or(state_diff) {
None => {
let n = selfcheck_n();
let pool = self.glm5_graph_pool(cache);
if let Some(st) = pool
.stages
.iter_mut()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
&& let Some(r) = st.runs.iter_mut().find(|r| r.lo == a && r.hi == b)
{
r.checked += 1;
}
note(&format!(
"self-check PASS dev={dev} stage=[{lo}, {hi}) run=[{a}, {b}) pos={pos} \
replay={k}/{n} phase={phase}: output ({width} elements) and the written \
recurrent state of {} layers are bit-identical to the eager walk on the \
same state{}",
b - a,
if k >= n {
"; this run is trusted for the session"
} else {
""
}
));
Ok((x_rep, true))
}
Some(what) => {
eprintln!(
"[glm5-decode-graph] SELF-CHECK FAILED dev={dev} stage=[{lo}, {hi}) \
run=[{a}, {b}) pos={pos} replay={k} phase={phase}: replay differs from the \
eager walk at {what}; the stage is latched EAGER for this session and this \
token is recomputed eager. The door did NOT engage."
);
restore(e, cache, &snaps)?;
let x_real =
self.hyper_range_decode_eager(e, topology, x_keep, a, b, pos_d, pos, cache)?;
self.glm5_graph_pool(cache).failed.push((dev, lo, hi));
Ok((x_real, false))
}
}
}
#[allow(clippy::too_many_arguments)]
#[allow(clippy::too_many_arguments)]
fn glm5_replay_run(
&self,
e: &Engine,
dev: usize,
lo: usize,
hi: usize,
a: usize,
b: usize,
x: CudaSlice<f32>,
width: usize,
pos_d: &CudaSlice<i32>,
cache: &mut Cache,
) -> Res<CudaSlice<f32>> {
fn ri_of(st: &StageGraphs, a: usize, b: usize) -> usize {
st.runs
.iter()
.position(|r| r.lo == a && r.hi == b)
.expect("run index resolved before the pieces loop")
}
let phase;
{
let pool = self.glm5_graph_pool(cache);
let st = pool
.stages
.iter_mut()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
.ok_or("glm5 decode graph: stage pool vanished between capture and replay")?;
let ri = st
.runs
.iter()
.position(|r| r.lo == a && r.hi == b)
.ok_or_else(|| format!("glm5 decode graph: no captured run [{a}, {b})"))?;
phase = st.runs[ri].phase;
let ctx = CapCtx {
dev,
lo,
hi,
run: ri,
runs: st.runs.len(),
a,
b,
phase,
recapture: false,
};
step(
e,
&ctx,
"memcpy_dtod(x -> x_io)",
e.copy_into(&mut st.x_io, 0, &x, width),
)?;
st.launched_since_sync = true;
}
let n_pieces = {
let pool = self.glm5_graph_pool(cache);
let st = pool
.stages
.iter()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
.expect("checked above");
st.runs[ri_of(st, a, b)].pieces.len()
};
for pi in 0..n_pieces {
let middle = {
let pool = self.glm5_graph_pool(cache);
let st = pool
.stages
.iter_mut()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
.expect("checked above");
let ri = ri_of(st, a, b);
let ctx = CapCtx {
dev,
lo,
hi,
run: ri,
runs: st.runs.len(),
a,
b,
phase,
recapture: false,
};
match &st.runs[ri].pieces[pi] {
RunPiece::Graph { graphs, .. } => {
let prev = e.f16_scratch_swap(st.f16.take());
let launched = graphs[phase].launch();
st.f16 = e.f16_scratch_swap(prev);
step(e, &ctx, "cuGraphLaunch", launched.map_err(Into::into))?;
None
}
RunPiece::Middle(il) => {
let il = *il;
let (s0, s1) = {
let [s0, s1] = &mut st.spare;
(
std::mem::replace(s0, e.zeros(1)?),
std::mem::replace(s1, e.zeros(1)?),
)
};
let h = std::mem::replace(&mut st.ws.h, s0);
let mixed = std::mem::replace(&mut st.mixed, s1);
Some((il, h, mixed))
}
}
};
if let Some((il, h, mut mixed)) = middle {
let prev = {
let pool = self.glm5_graph_pool(cache);
let st = pool
.stages
.iter_mut()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
.expect("checked above");
e.f16_scratch_swap(st.f16.take())
};
let r = self.hyper_mla_mid_post_ws(e, il, &h, pos_d, cache, &mut mixed);
let pool = self.glm5_graph_pool(cache);
let st = pool
.stages
.iter_mut()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
.expect("checked above");
st.f16 = e.f16_scratch_swap(prev);
st.spare[0] = std::mem::replace(&mut st.ws.h, h);
st.spare[1] = std::mem::replace(&mut st.mixed, mixed);
let ctx = CapCtx {
dev,
lo,
hi,
run: ri_of(st, a, b),
runs: st.runs.len(),
a,
b,
phase,
recapture: false,
};
step(e, &ctx, "mla_middle", r)?;
}
}
{
let pool = self.glm5_graph_pool(cache);
let st = pool
.stages
.iter_mut()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
.expect("checked above");
let ri = ri_of(st, a, b);
st.runs[ri].phase ^= 1;
}
for il in a..b {
if let Some(rl) = cache.recur[il].as_mut() {
std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
}
}
let out_ctx = CapCtx {
dev,
lo,
hi,
run: 0,
runs: 0,
a,
b,
phase: 0,
recapture: false,
};
let mut out = step(e, &out_ctx, "alloc(out)", e.uninit(width))?;
{
let pool = self.glm5_graph_pool(cache);
let st = pool
.stages
.iter()
.find(|s| s.dev == dev && s.lo == lo && s.hi == hi)
.expect("checked above");
step(
e,
&out_ctx,
"memcpy_dtod(x_out -> out)",
e.copy_into(&mut out, 0, &st.x_out, width),
)?;
}
crate::GLM5_DECODE_GRAPH_REPLAYS.fetch_add(1, Ordering::Relaxed);
Ok(out)
}
}
#[cfg(test)]
mod tests {
use super::capturable_runs;
#[test]
fn runs_split_on_the_uncapturable_layers() {
let cap = |il: usize| il % 4 != 3;
assert_eq!(
capturable_runs(0, 12, cap),
vec![(0, 3), (4, 7), (8, 11)],
"whole-trunk split"
);
assert_eq!(capturable_runs(5, 12, cap), vec![(5, 7), (8, 11)]);
assert_eq!(capturable_runs(0, 3, cap), vec![(0, 3)]);
assert!(capturable_runs(0, 4, |il| il == 4).is_empty());
assert!(capturable_runs(7, 7, |_| true).is_empty());
}
#[test]
fn every_run_length_makes_an_even_number_of_swaps() {
for (a, b) in capturable_runs(0, 12, |il| il % 4 != 3) {
assert_eq!(2 * (b - a) % 2, 0, "run [{a}, {b}) swaps must be even");
}
}
}