ferrox_models/
penalty_window.rs

1//! The one place that decides which tokens the repetition, presence
2//! and frequency penalties look back over.
3//!
4//! # What llama.cpp does, and where
5//!
6//! llama.cpp's penalties sampler is stateful: it keeps a ring buffer of
7//! the last `penalty_last_n` tokens it has ACCEPTED, plus a count map
8//! over that buffer, and `apply` walks the candidate list looking each
9//! candidate up in the map (`src/llama-sampler.cpp:2698-2759`). Nothing
10//! about that buffer knows whether a token was generated or read out of
11//! the prompt -- only that the sampler was told about it.
12//!
13//! Both front ends tell it about the prompt.
14//!
15//! - `llama-server` seeds the sampler with every prompt token before
16//!   the first token is drawn
17//!   (`tools/server/server-context.cpp:375-397`, the loop at 386-390:
18//!   `for (int i = 0; i < prompt.tokens.size(); i++) { ...
19//!   common_sampler_accept(smpl.get(), id, false); }`).
20//! - `llama-cli` does the same as it consumes the prompt, with the
21//!   reason written on the line above
22//!   (`tools/completion/completion.cpp:730-736`: *"push the prompt in
23//!   the sampling context in order to apply repetition penalties
24//!   later"*, `common_sampler_accept(smpl, embd_inp[n_consumed],
25//!   /* accept_grammar= */ false)`).
26//!
27//! `common_sampler_accept` pushes into the chain unconditionally
28//! (`common/sampling.cpp:472-504`), so a prompt token lands in the
29//! penalties ring buffer exactly like a generated one.
30//!
31//! So llama.cpp's window is the last `penalty_last_n` tokens of
32//! `prompt ++ generated`, and ferrox matches that. **This changes
33//! output** relative to ferrox before this module existed, on every run
34//! at the default `--repeat-penalty 1.1`: a token that occurs in the
35//! prompt is now penalised on its first generated occurrence.
36//!
37//! # Why it is a type and not a slice
38//!
39//! Because it was a slice, and five call sites each chose their own.
40//! `ferrox run`'s decode loops passed the generated tokens; the server's
41//! two decode loops passed the generated tokens (still do -- issue #73,
42//! the prompt ids do not reach that seam); `speculative` passed
43//! the prompt as well and then grew a `penalty_history_start` knob to
44//! paper over the disagreement; `draft_model` cloned the whole history
45//! per block; `kimi_generate` passed prompt and generated and was the
46//! only one that matched llama.cpp. Five sites, four answers, nothing
47//! enforcing agreement -- this repo's dominant bug shape.
48//!
49//! A [`PenaltyWindow`] is built from BOTH halves and there is no
50//! constructor that takes one slice, so a caller cannot produce a window
51//! without saying what its prompt is. A caller that genuinely has none
52//! writes `&[]` and that is visible in the diff.
53
54/// The tokens the penalties may see: a prompt and the tokens generated
55/// after it, in that order.
56///
57/// Borrowed rather than owned because this is built once per sampled
58/// token on every decode loop in the workspace; an owning window would
59/// clone the whole sequence per token.
60///
61/// The two halves are kept separate rather than concatenated because a
62/// decode loop already holds them separately, and concatenating would
63/// mean an allocation per token for a value only ever read back as "the
64/// last N of the two".
65#[derive(Debug, Clone, Copy)]
66pub struct PenaltyWindow<'a> {
67    prompt: &'a [usize],
68    generated: &'a [usize],
69}
70
71impl<'a> PenaltyWindow<'a> {
72    /// The window over `prompt` followed by `generated`.
73    ///
74    /// `prompt` is the tokens the model was fed before generation
75    /// started, and it belongs in the window: see the module docs for
76    /// the llama.cpp lines that put it there.
77    pub fn new(prompt: &'a [usize], generated: &'a [usize]) -> Self {
78        PenaltyWindow { prompt, generated }
79    }
80
81    /// Total tokens in the sequence, before `penalty_last_n` truncates
82    /// it.
83    pub fn len(&self) -> usize {
84        self.prompt.len() + self.generated.len()
85    }
86
87    pub fn is_empty(&self) -> bool {
88        self.len() == 0
89    }
90
91    /// The most recent `last_n` tokens of `prompt ++ generated`, oldest
92    /// first.
93    ///
94    /// This is llama.cpp's ring buffer expressed as a view: the buffer
95    /// holds at most `penalty_last_n` entries and the oldest is dropped
96    /// on every accept (`src/llama-sampler.cpp:2707-2716`), so its
97    /// contents are exactly the tail of the accepted sequence.
98    ///
99    /// Order does not matter to any caller -- only the multiset does --
100    /// but it is the natural one anyway, and a test reads it.
101    pub fn recent(&self, last_n: usize) -> impl Iterator<Item = usize> + '_ {
102        let start = self.len().saturating_sub(last_n);
103        // Split the single cut point across the two halves. `start` is
104        // at most `len()`, so both indices are in range and neither
105        // subtraction can wrap.
106        let from_prompt = start.min(self.prompt.len());
107        let from_generated = start.saturating_sub(self.prompt.len());
108        self.prompt[from_prompt..]
109            .iter()
110            .chain(self.generated[from_generated..].iter())
111            .copied()
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    /// The window is the tail of `prompt ++ generated`, so it slides
120    /// across the seam between them rather than restarting at it.
121    ///
122    /// A window implemented as "the last N of `generated`, plus all of
123    /// `prompt`" would keep token 0 here, and a window implemented as
124    /// "the last N of `generated`" would keep neither prompt token.
125    /// llama.cpp's ring buffer keeps exactly the last N accepted tokens
126    /// whichever half they came from.
127    #[test]
128    fn the_window_is_the_tail_of_the_prompt_and_the_generation_together() {
129        let window = PenaltyWindow::new(&[0, 1, 2], &[3, 4]);
130        assert_eq!(window.len(), 5);
131        assert_eq!(window.recent(3).collect::<Vec<_>>(), vec![2, 3, 4]);
132        // The cut can land inside the prompt, inside the generation, or
133        // exactly on the seam.
134        assert_eq!(window.recent(4).collect::<Vec<_>>(), vec![1, 2, 3, 4]);
135        assert_eq!(window.recent(2).collect::<Vec<_>>(), vec![3, 4]);
136        // Wider than the sequence is the whole sequence, not a panic.
137        assert_eq!(window.recent(1000).collect::<Vec<_>>(), vec![0, 1, 2, 3, 4]);
138        assert_eq!(window.recent(0).count(), 0);
139    }
140
141    /// Nothing generated yet is still a non-empty window, which is the
142    /// whole point: the first sampled token is already penalised
143    /// against the prompt.
144    #[test]
145    fn a_prompt_alone_is_a_window() {
146        let window = PenaltyWindow::new(&[7, 7, 8], &[]);
147        assert!(!window.is_empty());
148        assert_eq!(window.recent(64).collect::<Vec<_>>(), vec![7, 7, 8]);
149        assert_eq!(window.recent(2).collect::<Vec<_>>(), vec![7, 8]);
150    }
151
152    /// And an empty prompt is not a special case.
153    #[test]
154    fn an_empty_prompt_leaves_the_generated_tail() {
155        let window = PenaltyWindow::new(&[], &[1, 2, 3]);
156        assert_eq!(window.recent(2).collect::<Vec<_>>(), vec![2, 3]);
157        assert!(PenaltyWindow::new(&[], &[]).is_empty());
158    }
159}
ferrox_models/penalty_window.rs

ferrox_models/
penalty_window.rs
