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
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
//! Verify m-scaling probe (35B spec-gap triage, lane/close35 2026-07-08): times the spec
//! VERIFY forward (`decode_step_t_h_emb_dev` — the exact hot-loop kernel chain: resident-embed
//! gather, device logits, no host logits dtoh) at m = 1,2,3,4,6 from a FIXED primed depth,
//! rolling the cache back between calls so every timed call sees the identical state.
//! Prints us/call median + p10/p90 per m, the m=1-normalized cost curve, and the eager
//! decode_step reference. Read the curve against llama-bench `-d <depth> -p 1,2,3,4,6 -n 0`
//! (their verify batch = llama_decode of m tokens at depth — same dispatch as their MTP verify).
//!
//! MEASUREMENT-ONLY: no kernel/dispatch change; pure Instant+sync timing around existing calls.
//!
//! usage: verify-mscale <model.gguf> [depth=512] [reps=40] [m-list="1,2,3,4,6"]
//! env: fast-path core is default-on; nothing to set (MoE cache included).
//! `MEMRA_MSCALE_INTERLEAVE=1` alternates the m-list order on each repetition so width
//! comparisons share the same thermal/clock regime.
use memra_engine::Engine;
use memra_engine::forward::argmax;
use memra_engine::hybrid::HybridModel;
use memra_gguf::GgufFile;
fn pct(sorted: &[f64], p: f64) -> f64 {
let i = ((sorted.len() as f64 - 1.0) * p).round() as usize;
sorted[i]
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let path = std::env::args()
.nth(1)
.expect("usage: verify-mscale <model> [depth] [reps] [ms]");
let depth: usize = std::env::args()
.nth(2)
.and_then(|s| s.parse().ok())
.unwrap_or(512);
let reps: usize = std::env::args()
.nth(3)
.and_then(|s| s.parse().ok())
.unwrap_or(40);
let ms: Vec<usize> = std::env::args()
.nth(4)
.unwrap_or_else(|| "1,2,3,4,6".to_string())
.split(',')
.filter_map(|s| s.parse().ok())
.collect();
let max_m = ms.iter().copied().max().unwrap_or(6);
let e = Engine::new(0)?;
let g = GgufFile::open(&path)?;
let model = HybridModel::load(&e, &g)?;
let n_embd = model.cfg.n_embd as usize;
// Same synthetic prompt family as decode-bench (comparable depth state).
let prompt: Vec<u32> = (0..depth).map(|i| (100 + (i * 7) % 900) as u32).collect();
// Use the same stage-owned cache allocator as serving when the PP door is open. A
// primary-device Cache makes remote stages peer-read their KV and turns this measurement
// into the pre-ppspec placement bug rather than the live verify path.
let mut cache = memra_engine::pp::new_cache(&e, &model.cfg, depth + max_m + 32)?;
let t_prime = std::time::Instant::now();
let mut last_logits: Vec<f32> = if depth >= memra_engine::hybrid_forward::PRIME_MIN_T {
let (l, _h, _hiddens) = model.prime_cache(&e, &prompt, &mut cache, 0)?;
l
} else {
let mut l = Vec::new();
for &t in &prompt {
l = model.decode_step(&e, t, &mut cache)?;
}
l
};
e.stream().synchronize()?;
println!(
"primed depth={} in {:.2}s",
cache.pos,
t_prime.elapsed().as_secs_f64()
);
// Realistic verify tokens: the model's OWN greedy continuation (real expert routing),
// generated eagerly then rolled back.
let snap = cache.snapshot(&e)?;
let pos0 = cache.pos;
let mut gold: Vec<u32> = Vec::with_capacity(max_m);
let mut ll = last_logits.clone();
for _ in 0..max_m {
let nx = argmax(&ll) as u32;
gold.push(nx);
ll = model.decode_step(&e, nx, &mut cache)?;
}
cache.rollback(&e, &snap, 0)?;
let _ = &mut last_logits;
println!("verify tokens (greedy continuation): {gold:?}");
// Resident embed table — the spec hot loop's gather source.
let embd_gpu = model
.embd_gpu
.get_or_init(|| e.upload_u8(&model.embd.raw).expect("embed table upload"));
let (embd_qt, embd_rb) = model.embd.qt_and_row_bytes(n_embd);
// MEMRA_MSCALE_NOEAGER=1: skip the eager reference (keeps an nsys trace verify-only).
// MEMRA_MSCALE_PROFILE=1: bracket the TIMED verify reps in cuProfilerStart/Stop so
// `nsys --capture-range=cudaProfilerApi` records ONLY the verify kernel chain (run with a
// single-m list; warmups + rollbacks sit outside the bracket per rep is not possible —
// rollback D2D copies are inside the window and must be subtracted by name).
let profile = std::env::var("MEMRA_MSCALE_PROFILE").is_ok();
// Eager decode_step reference (the plain-decode per-token cost at this depth).
if std::env::var("MEMRA_MSCALE_NOEAGER").is_err() {
for _ in 0..3 {
let _ = model.decode_step(&e, gold[0], &mut cache)?;
cache.rollback(&e, &snap, 0)?;
}
e.stream().synchronize()?;
let mut ts: Vec<f64> = Vec::with_capacity(reps);
for _ in 0..reps {
e.stream().synchronize()?;
let t0 = std::time::Instant::now();
let _ = model.decode_step(&e, gold[0], &mut cache)?;
e.stream().synchronize()?;
ts.push(t0.elapsed().as_secs_f64() * 1e6);
cache.rollback(&e, &snap, 0)?;
}
ts.sort_by(|a, b| a.partial_cmp(b).unwrap());
println!(
"eager decode_step @d{depth}: median {:8.1} us p10 {:8.1} p90 {:8.1}",
pct(&ts, 0.5),
pct(&ts, 0.1),
pct(&ts, 0.9)
);
}
let interleave = std::env::var("MEMRA_MSCALE_INTERLEAVE").as_deref() == Ok("1");
let mut samples: Vec<Vec<f64>> = ms.iter().map(|_| Vec::with_capacity(reps)).collect();
if interleave {
// Warm every width before collecting any scored sample, then alternate
// forward/reverse order per repetition.
for &m in &ms {
let toks = &gold[0..m];
for _ in 0..3 {
let _ = model.decode_step_t_h_emb_dev(
&e,
toks,
pos0,
&mut cache,
Some((embd_gpu, embd_qt, embd_rb)),
)?;
cache.rollback(&e, &snap, 0)?;
}
}
e.stream().synchronize()?;
if profile {
unsafe {
cudarc::driver::sys::cuProfilerStart().result()?;
}
}
println!("measurement order: alternating forward/reverse by repetition");
for rep in 0..reps {
for ordinal in 0..ms.len() {
let i = if rep % 2 == 0 {
ordinal
} else {
ms.len() - 1 - ordinal
};
let m = ms[i];
let toks = &gold[0..m];
e.stream().synchronize()?;
let t0 = std::time::Instant::now();
let _ = model.decode_step_t_h_emb_dev(
&e,
toks,
pos0,
&mut cache,
Some((embd_gpu, embd_qt, embd_rb)),
)?;
e.stream().synchronize()?;
samples[i].push(t0.elapsed().as_secs_f64() * 1e6);
cache.rollback(&e, &snap, 0)?;
}
}
if profile {
unsafe {
cudarc::driver::sys::cuProfilerStop().result()?;
}
}
} else {
// Preserve the original width-major tool behavior when the research door is unset:
// warm, profile and score each width as one block.
for (i, &m) in ms.iter().enumerate() {
let toks = &gold[0..m];
for _ in 0..3 {
let _ = model.decode_step_t_h_emb_dev(
&e,
toks,
pos0,
&mut cache,
Some((embd_gpu, embd_qt, embd_rb)),
)?;
cache.rollback(&e, &snap, 0)?;
}
e.stream().synchronize()?;
if profile {
unsafe {
cudarc::driver::sys::cuProfilerStart().result()?;
}
}
for _ in 0..reps {
e.stream().synchronize()?;
let t0 = std::time::Instant::now();
let _ = model.decode_step_t_h_emb_dev(
&e,
toks,
pos0,
&mut cache,
Some((embd_gpu, embd_qt, embd_rb)),
)?;
e.stream().synchronize()?;
samples[i].push(t0.elapsed().as_secs_f64() * 1e6);
cache.rollback(&e, &snap, 0)?;
}
if profile {
unsafe {
cudarc::driver::sys::cuProfilerStop().result()?;
}
}
}
}
let mut med1 = 0.0f64;
for (i, &m) in ms.iter().enumerate() {
let ts = &mut samples[i];
ts.sort_by(|a, b| a.partial_cmp(b).unwrap());
let med = pct(ts, 0.5);
if i == 0 {
med1 = med;
}
println!(
"verify m={m} @d{depth}: median {:8.1} us p10 {:8.1} p90 {:8.1} | x{:.3} vs m={} | {:7.1} us/tok",
med,
pct(ts, 0.1),
pct(ts, 0.9),
med / med1,
ms[0],
med / m as f64
);
}
Ok(())
}