1515//! so a drafter changes the acceptance rate, never the output:
1616//!
1717//! - **prompt-lookup** ([`PromptLookupDrafter`]): the next tokens are guessed by
18- //! finding the most recent earlier occurrence of the current suffix in the
19- //! realized sequence and copying what followed it. No draft model, no
20- //! training, no weights — parametric over any model and any input. It shines
18+ //! finding the most recent earlier occurrence of the LONGEST recurring suffix
19+ //! of the realized sequence and copying what followed it. No draft model, no
20+ //! training, no weights, and no n-gram cap — parametric over any model and any
21+ //! input, at O(n) per draft. It shines
2122//! on structured/repetitive text (code, JSON, retrieval, format echoing) and
2223//! returns nothing on novel text, so the caller falls back to one plain step.
2324//! - **draft model** ([`ModelDrafter`]): a small second model (sharing the
@@ -92,14 +93,14 @@ pub fn draft_pairing_refusal(
9293}
9394
9495/// The zero-weight prompt-lookup drafter (the shipped default): stateless — it
95- /// reads the realized sequence and needs no prefill or commit.
96- pub struct PromptLookupDrafter {
97- pub ngram_max : usize ,
98- }
96+ /// reads the realized sequence and needs no prefill or commit. It carries no
97+ /// tuning parameter: the context it drafts from is the LONGEST recurrence the
98+ /// sequence itself contains.
99+ pub struct PromptLookupDrafter ;
99100
100101impl Drafter for PromptLookupDrafter {
101102 fn propose ( & mut self , realized : & [ i64 ] , cap : usize ) -> Result < Vec < i64 > > {
102- Ok ( prompt_lookup_draft ( realized, self . ngram_max , cap) )
103+ Ok ( prompt_lookup_draft ( realized, cap) )
103104 }
104105 fn commit ( & mut self , _accepted : usize , _bonus : i64 ) -> Result < ( ) > {
105106 Ok ( ( ) )
@@ -187,35 +188,76 @@ fn argmax(row: &[f32]) -> u32 {
187188
188189/// Draft up to `max_draft` continuation tokens for `seq` by prompt-lookup.
189190///
190- /// For the largest suffix length `g ∈ [1, ngram_max]` that occurs earlier in
191- /// `seq`, take the most recent such earlier occurrence and copy the up-to-
192- /// `max_draft` tokens that followed it. A longer matched suffix is a more
193- /// specific context, so it is tried first (higher acceptance). Returns an empty
194- /// draft when nothing matches — the signal to fall back to a single decode step.
191+ /// The drafting context is the LONGEST suffix of `seq` that occurs earlier in
192+ /// `seq` at all — there is no n-gram cap, because a cap is a guess about how
193+ /// much context an arbitrary input recurs over. Among earlier occurrences of
194+ /// that longest suffix the MOST RECENT is used (the freshest continuation), and
195+ /// the up-to-`max_draft` tokens that followed it are the draft. Returns an empty
196+ /// draft when the suffix never recurs — the signal to fall back to a single
197+ /// decode step.
198+ ///
199+ /// Found in O(n) time and space with the Z-function of the REVERSED sequence:
200+ /// `z[i]` is the longest common prefix of `rev` and `rev[i..]`, i.e. the longest
201+ /// common SUFFIX of `seq` and `seq[..n-i]` — an earlier occurrence ending at
202+ /// `n - i`. Scanning `i` ascending and keeping strictly longer matches yields the
203+ /// longest suffix, tie-broken toward the most recent occurrence. (The former
204+ /// capped scan was `O(n² · g)`, which an arbitrarily long realized sequence
205+ /// cannot afford.)
195206///
196207/// This never affects correctness, only speed: every drafted token is verified
197208/// and a wrong one is rejected. It is a pure function of the realized sequence.
198- pub fn prompt_lookup_draft ( seq : & [ i64 ] , ngram_max : usize , max_draft : usize ) -> Vec < i64 > {
209+ pub fn prompt_lookup_draft ( seq : & [ i64 ] , max_draft : usize ) -> Vec < i64 > {
199210 let n = seq. len ( ) ;
200- if n < 2 || max_draft == 0 || ngram_max == 0 {
211+ if n < 2 || max_draft == 0 {
201212 return Vec :: new ( ) ;
202213 }
203- // Longest matching suffix first: a more specific context accepts better.
204- let g_max = ngram_max. min ( n - 1 ) ;
205- for g in ( 1 ..=g_max) . rev ( ) {
206- let needle = & seq[ n - g..] ;
207- // Most recent earlier occurrence first (skip the trailing suffix itself).
208- for start in ( 0 ..n - g) . rev ( ) {
209- if & seq[ start..start + g] == needle {
210- let from = start + g;
211- let take = max_draft. min ( n - from) ;
212- if take > 0 {
213- return seq[ from..from + take] . to_vec ( ) ;
214- }
215- }
214+ let rev: Vec < i64 > = seq. iter ( ) . rev ( ) . copied ( ) . collect ( ) ;
215+ let z = z_function ( & rev) ;
216+
217+ // `best_end` is the EXCLUSIVE end of the earlier occurrence of the longest
218+ // recurring suffix; the draft is what followed it.
219+ let ( mut best_len, mut best_end) = ( 0usize , 0usize ) ;
220+ for ( i, & zi) in z. iter ( ) . enumerate ( ) . skip ( 1 ) {
221+ // The occurrence must fit entirely before the position it ends at.
222+ let len = zi. min ( n - i) ;
223+ if len > best_len {
224+ best_len = len;
225+ best_end = n - i;
226+ }
227+ }
228+ if best_len == 0 {
229+ return Vec :: new ( ) ;
230+ }
231+ let take = max_draft. min ( n - best_end) ;
232+ if take == 0 {
233+ return Vec :: new ( ) ;
234+ }
235+ seq[ best_end..best_end + take] . to_vec ( )
236+ }
237+
238+ /// `z[i]` = length of the longest common prefix of `s` and `s[i..]` (`z[0] = n`).
239+ /// Linear time, the standard two-pointer construction.
240+ fn z_function ( s : & [ i64 ] ) -> Vec < usize > {
241+ let n = s. len ( ) ;
242+ let mut z = vec ! [ 0usize ; n] ;
243+ if n == 0 {
244+ return z;
245+ }
246+ z[ 0 ] = n;
247+ let ( mut l, mut r) = ( 0usize , 0usize ) ;
248+ for i in 1 ..n {
249+ if i < r {
250+ z[ i] = ( r - i) . min ( z[ i - l] ) ;
251+ }
252+ while i + z[ i] < n && s[ z[ i] ] == s[ i + z[ i] ] {
253+ z[ i] += 1 ;
254+ }
255+ if i + z[ i] > r {
256+ l = i;
257+ r = i + z[ i] ;
216258 }
217259 }
218- Vec :: new ( )
260+ z
219261}
220262
221263#[ cfg( test) ]
@@ -224,20 +266,18 @@ mod tests {
224266
225267 #[ test]
226268 fn empty_or_trivial_sequence_drafts_nothing ( ) {
227- assert ! ( prompt_lookup_draft( & [ ] , 3 , 4 ) . is_empty( ) ) ;
228- assert ! ( prompt_lookup_draft( & [ 7 ] , 3 , 4 ) . is_empty( ) ) ;
229- // A budget or ngram of zero disables drafting.
230- assert ! ( prompt_lookup_draft( & [ 1 , 2 , 3 ] , 0 , 4 ) . is_empty( ) ) ;
231- assert ! ( prompt_lookup_draft( & [ 1 , 2 , 3 ] , 3 , 0 ) . is_empty( ) ) ;
269+ assert ! ( prompt_lookup_draft( & [ ] , 4 ) . is_empty( ) ) ;
270+ assert ! ( prompt_lookup_draft( & [ 7 ] , 4 ) . is_empty( ) ) ;
271+ // A zero draft budget disables drafting.
272+ assert ! ( prompt_lookup_draft( & [ 1 , 2 , 3 ] , 0 ) . is_empty( ) ) ;
232273 }
233274
234275 #[ test]
235276 fn repeated_suffix_drafts_its_earlier_continuation ( ) {
236277 // "a b c d ... a b" → the suffix `a b` last continued with `c d e`.
237278 let seq = [ 10 , 20 , 30 , 40 , 99 , 10 , 20 ] ;
238- let draft = prompt_lookup_draft ( & seq, 3 , 3 ) ;
239279 assert_eq ! (
240- draft ,
280+ prompt_lookup_draft ( & seq , 3 ) ,
241281 vec![ 30 , 40 , 99 ] ,
242282 "drafts what followed the earlier `10,20`"
243283 ) ;
@@ -246,77 +286,64 @@ mod tests {
246286 #[ test]
247287 fn longest_suffix_wins_over_a_shorter_ambiguous_one ( ) {
248288 // The 2-gram `2 3` occurred once (→ 4); the 1-gram `3` also occurred
249- // earlier (→ 4 as well) . The longer, more specific match is used.
289+ // earlier. The longer, more specific match is used.
250290 let seq = [ 1 , 2 , 3 , 4 , 5 , 2 , 3 ] ;
251- let draft = prompt_lookup_draft ( & seq, 3 , 2 ) ;
252- assert_eq ! ( draft, vec![ 4 , 5 ] , "the 2-gram `2,3` context drafts `4,5`" ) ;
291+ assert_eq ! ( prompt_lookup_draft( & seq, 2 ) , vec![ 4 , 5 ] ) ;
253292 }
254293
255294 #[ test]
256295 fn most_recent_occurrence_is_preferred ( ) {
257- // `7` appears after position 0 (→ 1) and position 3 (→ 8). The most
258- // recent earlier occurrence (index 3) drafts its follower .
296+ // `7` appears after position 0 (→ 1) and position 3 (→ 8). Among equally
297+ // long matches the freshest is used .
259298 let seq = [ 7 , 1 , 2 , 7 , 8 , 9 , 7 ] ;
260- let draft = prompt_lookup_draft ( & seq, 1 , 1 ) ;
261- assert_eq ! ( draft, vec![ 8 ] , "the freshest match for `7` drafts `8`" ) ;
299+ assert_eq ! ( prompt_lookup_draft( & seq, 1 ) , vec![ 8 ] ) ;
262300 }
263301
264302 #[ test]
265303 fn max_draft_caps_the_length ( ) {
266304 let seq = [ 5 , 6 , 7 , 8 , 9 , 5 , 6 ] ;
267- assert_eq ! ( prompt_lookup_draft( & seq, 3 , 2 ) , vec![ 7 , 8 ] ) ;
268- assert_eq ! ( prompt_lookup_draft( & seq, 3 , 10 ) , vec![ 7 , 8 , 9 , 5 , 6 ] ) ;
305+ assert_eq ! ( prompt_lookup_draft( & seq, 2 ) , vec![ 7 , 8 ] ) ;
306+ assert_eq ! ( prompt_lookup_draft( & seq, 10 ) , vec![ 7 , 8 , 9 , 5 , 6 ] ) ;
269307 }
270308
271309 #[ test]
272310 fn novel_suffix_drafts_nothing ( ) {
273311 // The trailing `42` never occurred earlier → no draft, fall back to a step.
274312 let seq = [ 1 , 2 , 3 , 4 , 42 ] ;
275- assert ! ( prompt_lookup_draft( & seq, 3 , 4 ) . is_empty( ) ) ;
313+ assert ! ( prompt_lookup_draft( & seq, 4 ) . is_empty( ) ) ;
276314 }
277315
278- // ── draft-pairing compatibility policy (row `speculative-draft-pairing`) ──
279-
280316 #[ test]
281- fn an_equal_or_larger_draft_is_compatible ( ) {
282- // Same-family self-pairing: identical vocab + context — the guaranteed
283- // case the browser witness uses.
284- assert ! ( draft_pairing_refusal ( 512 , 128 , 512 , 128 ) . is_none ( ) ) ;
285- // A draft that COVERS the target (larger vocab, longer context) is fine —
286- // every target id indexes the draft and the sequence never overflows it.
287- assert ! ( draft_pairing_refusal ( 512 , 128 , 1024 , 4096 ) . is_none ( ) ) ;
317+ fn a_longer_context_beats_a_more_recent_shorter_one ( ) {
318+ // The suffix `1,2,3` recurs (ending at 3) and continued with `7,8`.
319+ // The SHORTER suffix `2,3` has a MORE RECENT earlier occurrence (ending
320+ // at 7) that continued with `9`. Specificity wins: the longest recurring
321+ // context is the right one. A 2-token n-gram cap would have drafted `9`.
322+ let seq = [ 1 , 2 , 3 , 7 , 8 , 2 , 3 , 9 , 1 , 2 , 3 ] ;
323+ assert_eq ! ( prompt_lookup_draft ( & seq , 2 ) , vec! [ 7 , 8 ] ) ;
288324 }
289325
290326 #[ test]
291- fn a_narrower_vocabulary_is_refused ( ) {
292- // A target id ≥ the draft's vocab would index the draft's embedding out
293- // of range — refuse, naming both sizes.
294- let reason = draft_pairing_refusal ( 512 , 128 , 400 , 128 ) . expect ( "refused" ) ;
295- assert ! ( reason. contains( "400" ) && reason. contains( "512" ) , "{reason}" ) ;
296- assert ! ( reason. contains( "vocabulary" ) , "{reason}" ) ;
327+ fn the_drafter_has_no_context_cap_and_stays_linear ( ) {
328+ // An arbitrarily long recurrence is matched in full — no n-gram ceiling.
329+ // 4096 tokens is far past any cap a tuning constant would have imposed,
330+ // and the O(n) search returns immediately.
331+ let period: Vec < i64 > = ( 0 ..4096 ) . collect ( ) ;
332+ let mut seq = period. clone ( ) ;
333+ seq. push ( -1 ) ; // a separator so the recurrence is the whole period
334+ seq. extend_from_slice ( & period) ;
335+ // The longest recurring suffix is the entire 4096-token period, whose
336+ // earlier occurrence ended just before the separator.
337+ assert_eq ! ( prompt_lookup_draft( & seq, 1 ) , vec![ -1 ] ) ;
297338 }
298339
299340 #[ test]
300- fn a_shorter_context_is_refused ( ) {
301- // The target's realized sequence grows to 128; a 64-context draft would
302- // abort its forward when it crossed 64 — refuse before it can.
303- let reason = draft_pairing_refusal ( 512 , 128 , 512 , 64 ) . expect ( "refused" ) ;
304- assert ! ( reason. contains( "64" ) && reason. contains( "128" ) , "{reason}" ) ;
305- assert ! ( reason. contains( "context" ) , "{reason}" ) ;
306- }
307-
308- #[ test]
309- fn an_unknown_target_vocabulary_is_refused ( ) {
310- // Vocab 0 = the config declared none; coverage cannot be verified, so
311- // refuse rather than risk an out-of-range draft Gather.
312- assert ! ( draft_pairing_refusal( 0 , 128 , 512 , 128 ) . is_some( ) ) ;
313- }
314-
315- #[ test]
316- fn the_vocabulary_check_precedes_the_context_check ( ) {
317- // Both incompatible: the message names the vocabulary (the first gate),
318- // deterministically — a stable refusal reason, never order-dependent.
319- let reason = draft_pairing_refusal ( 512 , 128 , 400 , 64 ) . expect ( "refused" ) ;
320- assert ! ( reason. contains( "vocabulary" ) , "{reason}" ) ;
341+ fn z_function_matches_a_naive_reference ( ) {
342+ let s: Vec < i64 > = [ 1 , 2 , 1 , 2 , 1 , 3 , 1 , 2 , 1 ] . to_vec ( ) ;
343+ let z = z_function ( & s) ;
344+ for i in 0 ..s. len ( ) {
345+ let naive = ( 0 ..s. len ( ) - i) . take_while ( |& k| s[ k] == s[ i + k] ) . count ( ) ;
346+ assert_eq ! ( z[ i] , naive, "z[{i}]" ) ;
347+ }
321348 }
322349}
0 commit comments