memra_engine/spec_phase.rs
1//! Per-burst phase attribution for speculative rounds (`MEMRA_SPEC_TRACE`, generalized
2//! lane/glm5-extract-general from the glm5 loop's `MEMRA_GLM5_SPEC_TRACE` — the alias
3//! stays honored). The draft / verify / accept / rollback / source-maintenance split is
4//! SPEC-FAMILY-GENERIC: any spec loop owns those five boundaries, and the level-2 verify
5//! sub-split buckets are MIXER-CLASS buckets (KDA, MLA — multi-family classes), not one
6//! model's. The emit TAGS are the caller's, so a family's banked receipts keep their
7//! exact grep shape (`[glm5-phase]` / `[glm5-phase-v]` for the glm5 loop).
8//!
9//! DEFAULT OFF BY DESIGN (the flag row's law): each phase boundary SYNCHRONIZES the
10//! stream so device time lands in the right bucket, which serializes the round — a
11//! diagnostic instrument, never a serving mode, and its numbers are phase SHARES, not
12//! round walls (the un-traced round overlaps what the trace separates).
13
14use crate::Engine;
15use std::sync::atomic::Ordering;
16
17/// `MEMRA_SPEC_TRACE=1` (or the glm5 alias): per-burst phase attribution is on.
18pub fn spec_trace_on() -> bool {
19 spec_trace_level() >= 1
20}
21
22/// Trace LEVEL: `1` = the per-burst phase lines (draft/verify/accept/roll/maint);
23/// `2` = additionally the VERIFY sub-split — batched-class vs sequential-class time per
24/// burst (vkda with its in-kernel scan share, vmla, vrest = glue+FFN+head). Level 2 adds
25/// per-layer stream drains on top of level 1's phase drains: shares, never walls, never
26/// a perf row (the standing trace law). Read once per process (the worker chunk-policy
27/// pattern). The general name wins when both names are set to DIFFERENT levels — with
28/// one loud stderr line naming the override (the alias is never silently dead).
29pub fn spec_trace_level() -> u8 {
30 use std::sync::OnceLock;
31 static L: OnceLock<u8> = OnceLock::new();
32 *L.get_or_init(|| {
33 spec_trace_level_from(
34 std::env::var("MEMRA_SPEC_TRACE").ok().as_deref(),
35 std::env::var("MEMRA_GLM5_SPEC_TRACE").ok().as_deref(),
36 )
37 })
38}
39
40fn parse_level(v: Option<&str>) -> Option<u8> {
41 match v {
42 Some("1") => Some(1),
43 Some("2") => Some(2),
44 _ => None,
45 }
46}
47
48/// Pure resolution over the general name and the glm5 alias (unit-tested without env
49/// mutation). Either name alone is honored; both set and disagreeing = the general name
50/// wins LOUDLY (one stderr line naming both values).
51fn spec_trace_level_from(general: Option<&str>, glm5_alias: Option<&str>) -> u8 {
52 let g = parse_level(general);
53 let a = parse_level(glm5_alias);
54 if let (Some(gv), Some(av)) = (g, a)
55 && gv != av
56 {
57 eprintln!(
58 "[spec-trace] MEMRA_SPEC_TRACE={gv} overrides MEMRA_GLM5_SPEC_TRACE={av} \
59 (the general flag wins; unset one to silence this)"
60 );
61 }
62 g.or(a).unwrap_or(0)
63}
64
65/// Trace-level-2 verify sub-phase accumulators (ns), drained by [`SpecPhaseNs::emit`].
66/// Module-level atomics so the walk needs no signature plumbing through the ppN twin
67/// (the KDA_FUSED6_DISPATCHES precedent); level 2 is a single-session instrument, so
68/// cross-session interleaving is out of scope by definition.
69pub(crate) static V_KDA_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
70pub(crate) static V_KDA_SCAN_NS: std::sync::atomic::AtomicU64 =
71 std::sync::atomic::AtomicU64::new(0);
72pub(crate) static V_MLA_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
73/// FFN-branch share of the verify walk (lane/glm5-vrest): MoE + dense + shexp time inside
74/// vrest, so the box window can split the vrest bucket without re-deriving it. Ticks only
75/// on the batched arm, like its siblings; vrest's own definition (verify - vkda - vmla)
76/// stays unchanged for cross-window comparability — the line prints vffn INSIDE vrest.
77pub(crate) static V_FFN_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
78
79/// Per-burst phase counters (ns) — the verify-toll dataset the dspark loop banks under
80/// its `stats` clocks (`ns_draft/ns_verify/...`, dflash.rs).
81#[derive(Default)]
82pub(crate) struct SpecPhaseNs {
83 pub(crate) draft: u64,
84 pub(crate) verify: u64,
85 pub(crate) accept: u64,
86 pub(crate) roll: u64,
87 pub(crate) maint: u64,
88 pub(crate) rounds: u64,
89}
90
91impl SpecPhaseNs {
92 /// Phase-boundary clock: drain the engines' streams so the elapsed time since the last
93 /// clock is attributable to the phase that just ran (the dspark `clock(stats, e)`
94 /// contract). `eh` == `e` when the ppN door is shut; under a split the verify walk's own
95 /// terminal drain already covers the stage streams transitively, so syncing the primary
96 /// and head streams here bounds every phase that runs on them.
97 pub(crate) fn clock(e: &Engine, eh: &Engine) -> std::time::Instant {
98 let _ = e.stream().synchronize();
99 if !std::ptr::eq(e, eh) {
100 let _ = eh.stream().synchronize();
101 }
102 std::time::Instant::now()
103 }
104
105 /// One line per burst under `tag`; the level-2 verify sub-split under `tag_v` — both
106 /// tags belong to the CALLING family so its banked receipts keep their grep shape.
107 pub(crate) fn emit(&self, tag: &str, tag_v: &str, k: usize) {
108 if self.rounds == 0 {
109 return;
110 }
111 let ms = |ns: u64| ns as f64 / 1e6;
112 let per = |ns: u64| ns as f64 / 1e6 / self.rounds as f64;
113 let total = self.draft + self.verify + self.accept + self.roll + self.maint;
114 eprintln!(
115 "[{tag}] rounds={} k={k} total={:.2}ms | draft={:.2} verify={:.2} \
116 accept={:.2} roll={:.2} maint={:.2} | per-round ms: draft={:.3} verify={:.3} \
117 accept={:.3} roll={:.3} maint={:.3} total={:.3}",
118 self.rounds,
119 ms(total),
120 ms(self.draft),
121 ms(self.verify),
122 ms(self.accept),
123 ms(self.roll),
124 ms(self.maint),
125 per(self.draft),
126 per(self.verify),
127 per(self.accept),
128 per(self.roll),
129 per(self.maint),
130 per(total),
131 );
132 // Level-2 verify sub-split (lane/glm5-verify-batch): batched-class vs
133 // sequential-class shares. vrest = the verify phase minus the mixer buckets
134 // (hc glue + FFN/MoE + head); scan = the sequential KDA chain inside the
135 // batched call. Drained per burst so consecutive bursts stay comparable.
136 if spec_trace_level() >= 2 {
137 let vkda = V_KDA_NS.swap(0, Ordering::Relaxed);
138 let scan = V_KDA_SCAN_NS.swap(0, Ordering::Relaxed);
139 let vmla = V_MLA_NS.swap(0, Ordering::Relaxed);
140 let vffn = V_FFN_NS.swap(0, Ordering::Relaxed);
141 let vrest = self.verify.saturating_sub(vkda + vmla);
142 eprintln!(
143 "[{tag_v}] rounds={} k={k} | per-round ms: vkda={:.3} (scan={:.3}) \
144 vmla={:.3} vrest={:.3} (vffn={:.3})",
145 self.rounds,
146 per(vkda),
147 per(scan),
148 per(vmla),
149 per(vrest),
150 per(vffn),
151 );
152 }
153 }
154}
155
156#[cfg(test)]
157mod tests {
158 use super::{parse_level, spec_trace_level_from};
159
160 #[test]
161 fn level_resolution_honors_both_names_general_wins() {
162 // off by default; junk values are off (the original match-arm law)
163 assert_eq!(spec_trace_level_from(None, None), 0);
164 assert_eq!(spec_trace_level_from(Some("x"), None), 0);
165 // either name alone
166 assert_eq!(spec_trace_level_from(Some("1"), None), 1);
167 assert_eq!(spec_trace_level_from(Some("2"), None), 2);
168 assert_eq!(spec_trace_level_from(None, Some("1")), 1);
169 assert_eq!(spec_trace_level_from(None, Some("2")), 2);
170 // agreement and (loud) general-wins disagreement
171 assert_eq!(spec_trace_level_from(Some("2"), Some("2")), 2);
172 assert_eq!(spec_trace_level_from(Some("1"), Some("2")), 1);
173 assert_eq!(spec_trace_level_from(Some("2"), Some("1")), 2);
174 // a junk general value never masks a valid alias
175 assert_eq!(spec_trace_level_from(Some("x"), Some("1")), 1);
176 assert_eq!(parse_level(Some("0")), None);
177 }
178}