Skip to main content

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}