memra_sampling/lib.rs
1//! Host-side sampler chain (BASE-2, MEMRA-BUILD-MAP §BASE-2). Ports llama.cpp CPU sampler
2//! semantics (llama-sampler.cpp): repetition/freq/presence penalties -> temperature -> top-k ->
3//! top-p -> min-p -> categorical draw. Greedy (temp<=0) = argmax, the bit-exact reference.
4//!
5//! Runs on the host over the full [n_vocab] f32 logit vector already brought back by the per-step
6//! D2H sync (decode.rs) — at B=2-4 this is single-µs, no GPU kernel needed (the GPU-fused sampler
7//! is a deferred PERF item, only needed once CUDA-graph removes the D2H barrier).
8
9use std::collections::HashMap;
10
11/// Sampler configuration. Defaults = greedy (temp 0). Order of application matches llama.cpp.
12#[derive(Clone, Debug)]
13pub struct SamplerConfig {
14 pub temperature: f32, // <= 0.0 => greedy argmax (penalties/top-k/p ignored)
15 pub top_k: usize, // 0 => disabled (keep all)
16 pub top_p: f32, // 1.0 => disabled
17 pub min_p: f32, // 0.0 => disabled
18 pub penalty_last_n: usize, // window of recent tokens for penalties (0 => disabled)
19 pub penalty_repeat: f32, // 1.0 => disabled (llama default 1.0)
20 pub penalty_freq: f32, // 0.0 => disabled
21 pub penalty_present: f32, // 0.0 => disabled
22 pub seed: u64,
23}
24
25impl Default for SamplerConfig {
26 fn default() -> Self {
27 SamplerConfig {
28 temperature: 0.0,
29 top_k: 0,
30 top_p: 1.0,
31 min_p: 0.0,
32 penalty_last_n: 0,
33 penalty_repeat: 1.0,
34 penalty_freq: 0.0,
35 penalty_present: 0.0,
36 seed: 0,
37 }
38 }
39}
40
41/// SESSION-RESUME SAMPLER IDENTITY (lane/session-resume-sampler-predicate-20260820; receipts
42/// `research/spec-cache-20260818/SESSION-RESUME-PREDICATE.md`).
43///
44/// The canonical form of a request's sampler, for exactly one question: **may this request resume
45/// a parked whole session that some OTHER request's sampler shaped?** The spec pool's resume probe
46/// compared prompts and never samplers — that omission is how a filtered request inherited a draft
47/// graph captured unfiltered (`memra-engine` `SampledGraphKey`, lane/graph-s-key-exactness-
48/// 20260819). Keying the graph closed the exactness hole; it did not make cross-sampler resume
49/// SOUND, and the house posture in that situation is refuse-on-ambiguity with the refusal naming
50/// itself. This type is that predicate.
51///
52/// CANONICALIZATION, and why each rule is safe. Two encodings that name the same program must
53/// compare equal, or the predicate refuses resumes that cost nothing to allow:
54/// - `temperature <= 0.0` is GREEDY — `-1.0` and `0.0` are one program, so `temp_bits` is pinned
55/// to `0.0` in that regime and `greedy` carries the distinction. The greedy/sampled flip is
56/// itself a refusal: the two arms consume a parked `next_pred`/`pending_tok` differently and
57/// engage different captured draft graphs.
58/// - `top_k == 0`, `top_p >= 1.0`, `min_p <= 0.0` are each the OFF sentinel (matching
59/// `is_spec_sampling` and `SampledGraphKey::pure_temp`), canonicalized so `top_p 1.5` and
60/// `top_p 1.0` do not look like a change.
61/// - Penalties are OFF as a group iff `penalty_last_n == 0` or all three coefficients are neutral
62/// — the same `pen_on` predicate `spec.rs` computes. Off canonicalizes to the whole disabled
63/// tuple, so `penalty_last_n 64` with neutral coefficients equals penalties absent.
64///
65/// Float fields compare by BITS after canonicalization (no NaN/`-0.0` surprise), the same
66/// discipline `SampledGraphKey` uses.
67///
68/// `seed` IS carried and DELIBERATELY NOT COMPARED — see [`SamplerIdentity::mismatch`].
69#[derive(Clone, Copy, PartialEq, Eq, Debug)]
70pub struct SamplerIdentity {
71 greedy: bool,
72 temp_bits: u32,
73 /// Carried for the record and for callers that want to log it; NOT part of `mismatch`.
74 seed: u64,
75 top_k: usize,
76 top_p_bits: u32,
77 min_p_bits: u32,
78 penalty_last_n: usize,
79 penalty_repeat_bits: u32,
80 penalty_freq_bits: u32,
81 penalty_present_bits: u32,
82}
83
84impl SamplerIdentity {
85 /// Canonical identity of a sampler configuration.
86 pub fn of(cfg: &SamplerConfig) -> Self {
87 let greedy = cfg.temperature <= 0.0;
88 let pen_on = cfg.penalty_last_n > 0
89 && (cfg.penalty_repeat != 1.0 || cfg.penalty_freq != 0.0 || cfg.penalty_present != 0.0);
90 SamplerIdentity {
91 greedy,
92 temp_bits: if greedy { 0.0f32 } else { cfg.temperature }.to_bits(),
93 seed: cfg.seed,
94 top_k: cfg.top_k,
95 top_p_bits: if cfg.top_p >= 1.0 { 1.0f32 } else { cfg.top_p }.to_bits(),
96 min_p_bits: if cfg.min_p <= 0.0 { 0.0f32 } else { cfg.min_p }.to_bits(),
97 penalty_last_n: if pen_on { cfg.penalty_last_n } else { 0 },
98 penalty_repeat_bits: if pen_on { cfg.penalty_repeat } else { 1.0f32 }.to_bits(),
99 penalty_freq_bits: if pen_on { cfg.penalty_freq } else { 0.0f32 }.to_bits(),
100 penalty_present_bits: if pen_on { cfg.penalty_present } else { 0.0f32 }.to_bits(),
101 }
102 }
103
104 /// The seed this identity was built from (logging/receipts only — never compared).
105 pub fn seed(&self) -> u64 {
106 self.seed
107 }
108
109 /// The FIRST field on which `self` (an incoming request) differs from `parked` (the sampler
110 /// that shaped a parked session), as a stable name for the refusal line — `None` when the two
111 /// samplers are equivalent and the resume is legal. A refusal that does not say why is
112 /// indistinguishable from an unwired mechanism, so the name is the deliverable, not a nicety.
113 ///
114 /// Order is fixed and coarsest-first (`regime` before the field that only exists inside one
115 /// regime), so the reported name is the most informative one rather than an artifact of struct
116 /// layout.
117 ///
118 /// **`seed` IS NOT COMPARED, deliberately.** It is the one sampler field a resume may change,
119 /// for two reasons that are both mechanical:
120 /// - The only parked state that BAKES the seed is the sampled draft graph, and
121 /// `SampledGraphKey` already carries `seed`: a seed change drops the parked graph and
122 /// recaptures. (`memra-engine` `spec.rs`; pinned by `seed_alone_still_rekeys_the_draft_graph`
123 /// in that crate's `sampled_graph_key` tests.)
124 /// - The session's persisted Philox counters (`SpecSession::sctr/uctr`) are counter-based:
125 /// `philox(seed', ctr)` continued from another seed's counter position is an independent
126 /// stream, not a repeated one. Reproducibility is already scoped per `(seed, session)`
127 /// rather than per seed (`memra-server` `worker.rs`, the spec-burst sampling note), so a
128 /// changed seed costs nothing that same-seed resume was not already costing.
129 ///
130 /// Comparing it would refuse essentially ALL sampled traffic: omitting `seed` on a serve
131 /// request draws fresh per-request entropy, so every turn of every seed-omitting conversation
132 /// would carry a "changed" seed. That is a cost with no soundness gain, which is exactly the
133 /// trade this predicate exists to make explicitly rather than by accident.
134 pub fn mismatch(&self, parked: &Self) -> Option<&'static str> {
135 if self.greedy != parked.greedy {
136 return Some("regime");
137 }
138 if self.temp_bits != parked.temp_bits {
139 return Some("temperature");
140 }
141 if self.top_k != parked.top_k {
142 return Some("top_k");
143 }
144 if self.top_p_bits != parked.top_p_bits {
145 return Some("top_p");
146 }
147 if self.min_p_bits != parked.min_p_bits {
148 return Some("min_p");
149 }
150 if self.penalty_last_n != parked.penalty_last_n {
151 return Some("penalty_last_n");
152 }
153 if self.penalty_repeat_bits != parked.penalty_repeat_bits {
154 return Some("penalty_repeat");
155 }
156 if self.penalty_freq_bits != parked.penalty_freq_bits {
157 return Some("penalty_freq");
158 }
159 if self.penalty_present_bits != parked.penalty_present_bits {
160 return Some("penalty_present");
161 }
162 None
163 }
164
165 /// THE PRE-LANE PREDICATE, RESTATED (teeth, not production). The spec pool-resume probe
166 /// applied no sampler test at all — it compared prompts and nothing else — so every sampler
167 /// pair was admitted. Restating it here keeps the refusal tests DECISIVE instead of
168 /// tautological: the same pair that `mismatch` names must be admitted by this, or the test is
169 /// asserting against a mechanism that never existed.
170 ///
171 /// It is also what `MEMRA_SPEC_RESUME_SAMPLER=0` selects at runtime (the rollback door and the
172 /// A/B arm the cost measurement needs), so this function is the single definition of "legacy"
173 /// for both the tests and the server.
174 pub fn legacy_admits(&self, _parked: &Self) -> bool {
175 true
176 }
177}
178
179/// Stateful sampler: owns the RNG + the recent-token history (for penalties).
180pub struct Sampler {
181 cfg: SamplerConfig,
182 rng: SplitMix64,
183 history: Vec<u32>, // recently emitted tokens (for penalty window)
184 // Counts over exactly the active penalty window. Keeping this incrementally makes the host
185 // oracle cheaper too, and lets the serving path upload O(unique ids) sparse penalty state
186 // instead of either the full vocabulary or an O(history^2) device-side dedup walk.
187 penalty_counts: HashMap<u32, u32>,
188}
189
190impl Sampler {
191 pub fn new(cfg: SamplerConfig) -> Self {
192 let rng = SplitMix64::new(cfg.seed);
193 Sampler {
194 cfg,
195 rng,
196 history: Vec::new(),
197 penalty_counts: HashMap::new(),
198 }
199 }
200
201 pub fn is_greedy(&self) -> bool {
202 self.cfg.temperature <= 0.0
203 }
204 /// Sampled spec in its FASTEST regime: pure temperature, no truncation filters, no
205 /// penalties. Filters and penalties are also distribution-exact under the rejection
206 /// verify (spec.rs applies both symmetrically to draft q and target p), so they remain
207 /// spec-ELIGIBLE — see `spec_eligible` in memra-server's worker, the authoritative
208 /// predicate. What they cost is the in-graph draft chain: the captured sampled draft
209 /// samples from the RAW softmax and can hold neither per-row filter stats nor a varying
210 /// penalty history, so `spec.rs` engages `graph_s` only in this pure-temp regime
211 /// (`pure_temp`) and otherwise falls back to the eager draft chain. This predicate names
212 /// that regime; it is NOT an eligibility test.
213 pub fn is_spec_sampling(&self) -> bool {
214 self.cfg.temperature > 0.0
215 && self.cfg.penalty_repeat == 1.0
216 && self.cfg.penalty_freq == 0.0
217 && self.cfg.penalty_present == 0.0
218 && self.cfg.top_k == 0
219 && self.cfg.top_p >= 1.0
220 && self.cfg.min_p <= 0.0
221 }
222 pub fn top_k(&self) -> usize {
223 self.cfg.top_k
224 }
225 pub fn penalty_last_n(&self) -> usize {
226 self.cfg.penalty_last_n
227 }
228 pub fn penalty_repeat(&self) -> f32 {
229 self.cfg.penalty_repeat
230 }
231 pub fn penalty_freq(&self) -> f32 {
232 self.cfg.penalty_freq
233 }
234 pub fn penalty_present(&self) -> f32 {
235 self.cfg.penalty_present
236 }
237 pub fn top_p(&self) -> f32 {
238 self.cfg.top_p
239 }
240 pub fn min_p(&self) -> f32 {
241 self.cfg.min_p
242 }
243 pub fn temperature(&self) -> f32 {
244 self.cfg.temperature
245 }
246 pub fn seed(&self) -> u64 {
247 self.cfg.seed
248 }
249 /// This sampler's canonical [`SamplerIdentity`] — the whole-session resume predicate's input.
250 pub fn identity(&self) -> SamplerIdentity {
251 SamplerIdentity::of(&self.cfg)
252 }
253
254 fn penalties_on(&self) -> bool {
255 self.cfg.penalty_last_n > 0
256 && (self.cfg.penalty_repeat != 1.0
257 || self.cfg.penalty_freq != 0.0
258 || self.cfg.penalty_present != 0.0)
259 }
260
261 /// Sparse `(token_id, count)` rows for the current penalty window. The order cannot affect
262 /// the arithmetic because every entry mutates one distinct logit; avoiding a per-token sort
263 /// is material on long agent histories.
264 pub fn penalty_counts(&self) -> Vec<(u32, u32)> {
265 debug_assert!(self.penalty_counts.values().all(|&count| count > 0));
266 self.penalty_counts
267 .iter()
268 .map(|(&id, &n)| (id, n))
269 .collect()
270 }
271
272 /// Record an emitted token so subsequent penalties see it.
273 pub fn accept(&mut self, token: u32) {
274 if self.penalties_on() {
275 let n = self.cfg.penalty_last_n;
276 if self.history.len() >= n {
277 let expired = self.history[self.history.len() - n];
278 let remove = {
279 let count = self
280 .penalty_counts
281 .get_mut(&expired)
282 .expect("active penalty window lost an accepted token");
283 *count -= 1;
284 *count == 0
285 };
286 if remove {
287 self.penalty_counts.remove(&expired);
288 }
289 }
290 *self.penalty_counts.entry(token).or_insert(0) += 1;
291 }
292 self.history.push(token);
293 }
294
295 /// Sample the next token id from raw logits [n_vocab]. Does NOT mutate logits in place beyond
296 /// a local copy. Returns the chosen token id. (Caller should `accept()` it afterwards.)
297 pub fn sample(&mut self, logits: &[f32]) -> u32 {
298 // Greedy fast path: argmax over RAW logits (penalties don't change the argmax direction
299 // enough to matter for the reference path; llama greedy is also pre-penalty argmax only
300 // when no penalties set — but to stay correct under penalties we still apply them first).
301 if self.is_greedy()
302 && self.cfg.penalty_repeat == 1.0
303 && self.cfg.penalty_freq == 0.0
304 && self.cfg.penalty_present == 0.0
305 {
306 return argmax_u32(logits);
307 }
308
309 // Work on (id, logit) candidates.
310 let mut cand: Vec<(u32, f32)> = logits
311 .iter()
312 .enumerate()
313 .map(|(i, &l)| (i as u32, l))
314 .collect();
315
316 // 1. Penalties (operate on logits, over the last-n history window).
317 self.apply_penalties(&mut cand);
318
319 // Greedy-with-penalties: argmax after penalties, no sampling.
320 if self.is_greedy() {
321 let mut best = cand[0];
322 for &c in &cand[1..] {
323 if c.1 > best.1 {
324 best = c;
325 }
326 }
327 return best.0;
328 }
329
330 // 2. Temperature scale.
331 if self.cfg.temperature > 0.0 && self.cfg.temperature != 1.0 {
332 let inv = 1.0 / self.cfg.temperature;
333 for c in cand.iter_mut() {
334 c.1 *= inv;
335 }
336 }
337
338 // 3. top-k: keep the k highest-logit candidates (partial sort by logit desc).
339 if self.cfg.top_k > 0 && self.cfg.top_k < cand.len() {
340 cand.sort_unstable_by(|a, b| b.1.total_cmp(&a.1));
341 cand.truncate(self.cfg.top_k);
342 }
343
344 // softmax over the surviving candidates (numerically stable).
345 softmax_inplace(&mut cand);
346
347 // 4. top-p (nucleus): smallest set whose cumulative prob >= top_p. Needs desc-by-prob order.
348 if self.cfg.top_p < 1.0 {
349 cand.sort_unstable_by(|a, b| b.1.total_cmp(&a.1));
350 let mut cum = 0.0f32;
351 let mut keep = 0usize;
352 for (i, c) in cand.iter().enumerate() {
353 cum += c.1;
354 keep = i + 1;
355 if cum >= self.cfg.top_p {
356 break;
357 }
358 }
359 cand.truncate(keep.max(1));
360 }
361
362 // 5. min-p: keep candidates with prob >= min_p * max_prob.
363 if self.cfg.min_p > 0.0 {
364 let maxp = cand.iter().map(|c| c.1).fold(0.0f32, f32::max);
365 let thresh = self.cfg.min_p * maxp;
366 cand.retain(|c| c.1 >= thresh);
367 if cand.is_empty() {
368 return argmax_u32(logits);
369 } // safety
370 }
371
372 // renormalize the surviving probs and draw.
373 let sum: f32 = cand.iter().map(|c| c.1).sum();
374 let r = self.rng.next_f32() * sum;
375 let mut acc = 0.0f32;
376 for c in &cand {
377 acc += c.1;
378 if acc >= r {
379 return c.0;
380 }
381 }
382 cand.last().unwrap().0
383 }
384
385 /// llama.cpp penalty: for each token in the last-n history, repeat-divide/multiply its logit
386 /// and apply frequency*count + presence. (llama-sampler.cpp penalties.)
387 fn apply_penalties(&self, cand: &mut [(u32, f32)]) {
388 let n = self.cfg.penalty_last_n;
389 if n == 0 {
390 return;
391 }
392 if self.cfg.penalty_repeat == 1.0
393 && self.cfg.penalty_freq == 0.0
394 && self.cfg.penalty_present == 0.0
395 {
396 return;
397 }
398 let start = self.history.len().saturating_sub(n);
399 let window = &self.history[start..];
400 if window.is_empty() {
401 return;
402 }
403 debug_assert_eq!(
404 self.penalty_counts
405 .values()
406 .map(|&n| n as usize)
407 .sum::<usize>(),
408 window.len(),
409 "incremental penalty counts must cover the active history window"
410 );
411 for c in cand.iter_mut() {
412 if let Some(&cnt) = self.penalty_counts.get(&c.0) {
413 // repeat: llama divides if logit>0 else multiplies (penalize toward 0)
414 if self.cfg.penalty_repeat != 1.0 {
415 if c.1 > 0.0 {
416 c.1 /= self.cfg.penalty_repeat;
417 } else {
418 c.1 *= self.cfg.penalty_repeat;
419 }
420 }
421 c.1 -= cnt as f32 * self.cfg.penalty_freq;
422 c.1 -= self.cfg.penalty_present; // presence: applied once if count>0
423 }
424 }
425 }
426}
427
428fn argmax_u32(logits: &[f32]) -> u32 {
429 let mut best = 0u32;
430 let mut bv = f32::NEG_INFINITY;
431 for (i, &v) in logits.iter().enumerate() {
432 if v > bv {
433 bv = v;
434 best = i as u32;
435 }
436 }
437 best
438}
439
440/// Stable softmax over candidate logits, writing probs back into the logit slot.
441fn softmax_inplace(cand: &mut [(u32, f32)]) {
442 let maxl = cand.iter().map(|c| c.1).fold(f32::NEG_INFINITY, f32::max);
443 let mut sum = 0.0f32;
444 for c in cand.iter_mut() {
445 let e = (c.1 - maxl).exp();
446 c.1 = e;
447 sum += e;
448 }
449 let inv = if sum > 0.0 { 1.0 / sum } else { 0.0 };
450 for c in cand.iter_mut() {
451 c.1 *= inv;
452 }
453}
454
455/// SplitMix64 — deterministic seedable RNG (so a fixed seed reproduces the token stream for the
456/// validation gate). Not crypto; fine for sampling.
457struct SplitMix64 {
458 state: u64,
459}
460impl SplitMix64 {
461 fn new(seed: u64) -> Self {
462 SplitMix64 {
463 state: seed.wrapping_add(0x9E3779B97F4A7C15),
464 }
465 }
466 fn next_u64(&mut self) -> u64 {
467 self.state = self.state.wrapping_add(0x9E3779B97F4A7C15);
468 let mut z = self.state;
469 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
470 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
471 z ^ (z >> 31)
472 }
473 /// uniform f32 in [0,1).
474 fn next_f32(&mut self) -> f32 {
475 // top 24 bits -> [0,1)
476 ((self.next_u64() >> 40) as f32) / (1u32 << 24) as f32
477 }
478}
479
480#[cfg(test)]
481mod tests {
482 use super::*;
483
484 fn sorted_counts(s: &Sampler) -> Vec<(u32, u32)> {
485 let mut counts = s.penalty_counts();
486 counts.sort_unstable_by_key(|&(id, _)| id);
487 counts
488 }
489
490 #[test]
491 fn greedy_is_argmax() {
492 let mut s = Sampler::new(SamplerConfig::default()); // temp 0
493 let logits = vec![0.1, 5.0, 2.0, -1.0];
494 assert_eq!(s.sample(&logits), 1);
495 }
496
497 #[test]
498 fn temp_sampling_deterministic_with_seed() {
499 let cfg = SamplerConfig {
500 temperature: 1.0,
501 seed: 42,
502 ..Default::default()
503 };
504 let logits = vec![1.0, 2.0, 3.0, 0.5];
505 let a = Sampler::new(cfg.clone()).sample(&logits);
506 let b = Sampler::new(cfg).sample(&logits);
507 assert_eq!(a, b, "same seed must reproduce the draw");
508 assert!(a < 4);
509 }
510
511 #[test]
512 fn top_k_one_is_argmax() {
513 let cfg = SamplerConfig {
514 temperature: 1.0,
515 top_k: 1,
516 seed: 7,
517 ..Default::default()
518 };
519 let logits = vec![0.1, 5.0, 2.0, -1.0];
520 assert_eq!(
521 Sampler::new(cfg).sample(&logits),
522 1,
523 "top_k=1 collapses to argmax"
524 );
525 }
526
527 #[test]
528 fn min_p_keeps_only_high_prob() {
529 // logit 10 dominates; min_p 0.5 should drop the rest -> always pick id 2.
530 let cfg = SamplerConfig {
531 temperature: 1.0,
532 min_p: 0.5,
533 seed: 3,
534 ..Default::default()
535 };
536 let logits = vec![0.0, 0.0, 10.0, 0.0];
537 for _ in 0..16 {
538 assert_eq!(Sampler::new(cfg.clone()).sample(&logits), 2);
539 }
540 }
541
542 #[test]
543 fn repeat_penalty_suppresses_recent() {
544 // greedy + heavy repeat penalty: id 1 is argmax but recently emitted -> should drop it.
545 let mut cfg = SamplerConfig::default();
546 cfg.penalty_last_n = 8;
547 cfg.penalty_repeat = 100.0;
548 let mut s = Sampler::new(cfg);
549 s.accept(1); // 1 was just emitted
550 let logits = vec![4.0, 5.0, 4.5, 1.0]; // raw argmax = 1
551 let got = s.sample(&logits);
552 assert_ne!(
553 got, 1,
554 "recent token must be penalized out of greedy argmax"
555 );
556 assert_eq!(got, 2, "next-highest after penalizing 1");
557 }
558
559 #[test]
560 fn penalty_counts_follow_the_sliding_window() {
561 let cfg = SamplerConfig {
562 penalty_last_n: 4,
563 penalty_repeat: 1.1,
564 penalty_freq: 0.5,
565 penalty_present: 0.25,
566 ..Default::default()
567 };
568 let mut s = Sampler::new(cfg);
569 for tok in [7, 8, 7, 9] {
570 s.accept(tok);
571 }
572 assert_eq!(sorted_counts(&s), vec![(7, 2), (8, 1), (9, 1)]);
573
574 s.accept(8); // active window is now [8, 7, 9, 8]
575 assert_eq!(sorted_counts(&s), vec![(7, 1), (8, 2), (9, 1)]);
576 s.accept(10); // active window is now [7, 9, 8, 10]
577 assert_eq!(sorted_counts(&s), vec![(7, 1), (8, 1), (9, 1), (10, 1)]);
578 s.accept(11); // active window is now [9, 8, 10, 11]; final 7 expires
579 assert_eq!(sorted_counts(&s), vec![(8, 1), (9, 1), (10, 1), (11, 1)]);
580 assert!(!s.penalty_counts().iter().any(|&(id, n)| id == 7 || n == 0));
581 }
582
583 #[test]
584 fn full_context_and_disabled_penalty_counts_are_exact() {
585 let mut full = Sampler::new(SamplerConfig {
586 penalty_last_n: usize::MAX,
587 penalty_present: 1.5,
588 ..Default::default()
589 });
590 for tok in [3, 3, 4, 5, 3] {
591 full.accept(tok);
592 }
593 assert_eq!(sorted_counts(&full), vec![(3, 3), (4, 1), (5, 1)]);
594
595 let mut neutral = Sampler::new(SamplerConfig {
596 penalty_last_n: usize::MAX,
597 ..Default::default()
598 });
599 neutral.accept(3);
600 assert!(neutral.penalty_counts().is_empty());
601 }
602}
603
604/// SESSION-RESUME SAMPLER PREDICATE teeth (lane/session-resume-sampler-predicate-20260820).
605/// CPU-only, no GPU: the predicate is a pure function, so its whole contract is testable here and
606/// a regression cannot hide behind "needs a card".
607///
608/// TEETH BOTH DIRECTIONS is the point. Every refusal test also asserts that `legacy_admits`
609/// ADMITS the same pair — the pre-lane probe compared prompts and never samplers — so the test is
610/// decisive (it fails on the old code) rather than tautological (passing because the pair was
611/// never resumable for some other reason).
612#[cfg(test)]
613mod resume_sampler_predicate_tests {
614 use super::*;
615
616 /// The vendor-default sampled shape the flip makes the majority of traffic.
617 fn vendor() -> SamplerConfig {
618 SamplerConfig {
619 temperature: 0.7,
620 top_k: 20,
621 top_p: 0.95,
622 seed: 20260820,
623 ..Default::default()
624 }
625 }
626
627 /// Today's pure-temp shape — the one that parks a `graph_s`.
628 fn pure_temp() -> SamplerConfig {
629 SamplerConfig {
630 temperature: 0.7,
631 seed: 20260820,
632 ..Default::default()
633 }
634 }
635
636 fn id(cfg: &SamplerConfig) -> SamplerIdentity {
637 SamplerIdentity::of(cfg)
638 }
639
640 // ---- direction 1: a SAME-sampler resume still resumes (no regression) ----
641
642 #[test]
643 fn identical_sampler_resumes() {
644 for cfg in [pure_temp(), vendor(), SamplerConfig::default()] {
645 assert_eq!(
646 id(&cfg).mismatch(&id(&cfg)),
647 None,
648 "a request must resume a session its own sampler shaped: {cfg:?}"
649 );
650 }
651 }
652
653 #[test]
654 fn disabled_sentinels_are_the_same_program() {
655 // top_p >= 1.0, min_p <= 0.0, top_k == 0 all mean OFF; a client that spells OFF
656 // differently on turn 2 must not lose its cache.
657 let a = SamplerConfig {
658 temperature: 0.7,
659 top_p: 1.0,
660 min_p: 0.0,
661 ..Default::default()
662 };
663 let b = SamplerConfig {
664 temperature: 0.7,
665 top_p: 1.5,
666 min_p: -1.0,
667 ..Default::default()
668 };
669 assert_eq!(id(&a).mismatch(&id(&b)), None, "off spelled two ways");
670 }
671
672 #[test]
673 fn greedy_temperature_encodings_are_one_program() {
674 let a = SamplerConfig {
675 temperature: 0.0,
676 ..Default::default()
677 };
678 let b = SamplerConfig {
679 temperature: -1.0,
680 ..Default::default()
681 };
682 assert_eq!(id(&a).mismatch(&id(&b)), None, "temp<=0 is one regime");
683 }
684
685 #[test]
686 fn neutral_penalty_coefficients_equal_penalties_absent() {
687 // penalty_last_n set but every coefficient neutral == `pen_on == false` in spec.rs.
688 let a = SamplerConfig {
689 temperature: 0.7,
690 penalty_last_n: 64,
691 penalty_repeat: 1.0,
692 penalty_freq: 0.0,
693 penalty_present: 0.0,
694 ..Default::default()
695 };
696 let b = SamplerConfig {
697 temperature: 0.7,
698 penalty_last_n: 0,
699 ..Default::default()
700 };
701 assert_eq!(
702 id(&a).mismatch(&id(&b)),
703 None,
704 "an inert penalty window is not a penalty change"
705 );
706 }
707
708 // ---- direction 2: a sampler-DIFFERING resume refuses, and names the field ----
709
710 #[test]
711 fn the_reproduced_collision_pair_refuses_and_names_a_filter() {
712 // The exact pair the predecessor reproduced on a live server: turn 1 pure-temp parks,
713 // turn 2 adds top_p 0.95 / top_k 20 and resumes. Same seed, same temperature.
714 let parked = id(&pure_temp());
715 let incoming = id(&vendor());
716 let field = incoming
717 .mismatch(&parked)
718 .expect("the reproduced collision pair must refuse");
719 assert_eq!(field, "top_k", "coarsest-first order names top_k here");
720 // DECISIVE: the pre-lane probe admitted exactly this pair.
721 assert!(
722 incoming.legacy_admits(&parked),
723 "legacy must admit the collision pair, or this test proves nothing"
724 );
725 }
726
727 #[test]
728 fn every_compared_field_refuses_on_its_own_and_names_itself() {
729 let base = pure_temp();
730 // A penalized base, so the three coefficients can each move ALONE: with penalties off on
731 // the parked side, turning any of them on also moves `penalty_last_n`, and coarsest-first
732 // order would (correctly) name the window instead of the coefficient.
733 let pen_base = SamplerConfig {
734 penalty_last_n: 64,
735 penalty_repeat: 1.1,
736 penalty_freq: 0.5,
737 penalty_present: 0.5,
738 ..base.clone()
739 };
740 // (field, parked, incoming) — exactly one canonical field differs in each row.
741 let cases: [(&str, SamplerConfig, SamplerConfig); 9] = [
742 (
743 "regime",
744 base.clone(),
745 SamplerConfig {
746 temperature: 0.0,
747 ..base.clone()
748 },
749 ),
750 (
751 "temperature",
752 base.clone(),
753 SamplerConfig {
754 temperature: 0.8,
755 ..base.clone()
756 },
757 ),
758 (
759 "top_k",
760 base.clone(),
761 SamplerConfig {
762 top_k: 20,
763 ..base.clone()
764 },
765 ),
766 (
767 "top_p",
768 base.clone(),
769 SamplerConfig {
770 top_p: 0.95,
771 ..base.clone()
772 },
773 ),
774 (
775 "min_p",
776 base.clone(),
777 SamplerConfig {
778 min_p: 0.05,
779 ..base.clone()
780 },
781 ),
782 (
783 "penalty_last_n",
784 pen_base.clone(),
785 SamplerConfig {
786 penalty_last_n: 128,
787 ..pen_base.clone()
788 },
789 ),
790 (
791 "penalty_repeat",
792 pen_base.clone(),
793 SamplerConfig {
794 penalty_repeat: 1.2,
795 ..pen_base.clone()
796 },
797 ),
798 (
799 "penalty_freq",
800 pen_base.clone(),
801 SamplerConfig {
802 penalty_freq: 0.6,
803 ..pen_base.clone()
804 },
805 ),
806 (
807 "penalty_present",
808 pen_base.clone(),
809 SamplerConfig {
810 penalty_present: 0.6,
811 ..pen_base.clone()
812 },
813 ),
814 ];
815 for (expect, parked_cfg, cfg) in cases {
816 let parked = id(&parked_cfg);
817 let incoming = id(&cfg);
818 assert_eq!(
819 incoming.mismatch(&parked),
820 Some(expect),
821 "changing {expect} alone must refuse and name {expect} ({cfg:?})"
822 );
823 assert!(
824 incoming.legacy_admits(&parked),
825 "legacy must admit the {expect} change, or the refusal test is tautological"
826 );
827 }
828 // Turning penalties ON from an unpenalized parked session is a `penalty_last_n` refusal —
829 // the coarsest true statement about that pair, asserted so the order is pinned.
830 assert_eq!(
831 id(&pen_base).mismatch(&id(&base)),
832 Some("penalty_last_n"),
833 "penalties on vs off is named at the window, not at a coefficient"
834 );
835 }
836
837 #[test]
838 fn greedy_to_sampled_and_back_both_refuse_as_regime() {
839 let g = id(&SamplerConfig::default());
840 let s = id(&pure_temp());
841 assert_eq!(s.mismatch(&g), Some("regime"));
842 assert_eq!(g.mismatch(&s), Some("regime"));
843 }
844
845 // ---- the seed decision, pinned so it cannot change silently ----
846
847 #[test]
848 fn seed_alone_does_not_refuse() {
849 // DELIBERATE (see SamplerIdentity::mismatch): the draft graph is re-keyed on seed by
850 // SampledGraphKey and the session's Philox counters are counter-based, so a changed seed
851 // is sound — and comparing it would refuse every seed-omitting sampled conversation,
852 // because an omitted seed draws fresh per-request entropy.
853 let a = pure_temp();
854 let b = SamplerConfig {
855 seed: 999,
856 ..a.clone()
857 };
858 assert_eq!(
859 id(&a).mismatch(&id(&b)),
860 None,
861 "seed is carried but not compared"
862 );
863 assert_ne!(id(&a).seed(), id(&b).seed(), "the seed is still recorded");
864 }
865
866 #[test]
867 fn mismatch_is_symmetric_and_identity_is_an_equivalence() {
868 let a = id(&pure_temp());
869 let b = id(&vendor());
870 assert_eq!(a.mismatch(&b).is_some(), b.mismatch(&a).is_some());
871 assert_eq!(a.mismatch(&a), None);
872 assert_eq!(b.mismatch(&b), None);
873 }
874}