Skip to main content

ferrox_models/
deepseek_v4_budget.rs

1//! DSV4 paged KV: sizing four heterogeneous tiers out of one budget, and
2//! the page-atomic allocator that hands out the window tier.
3//!
4//! A DSV4 layer stack does not have *a* KV cache. It has four pools with
5//! four different growth laws, and the whole difficulty is that they are
6//! bought with the same bytes:
7//!
8//! * the **window** tier -- the 128-position sliding KV ring, present on
9//!   every layer, page-granular;
10//! * the **compressed** tier -- one block per `ratio` full-history
11//!   tokens, on `ratio > 0` layers;
12//! * the **indexer** tier -- the Lightning-Indexer keys, one block per 4
13//!   full-history tokens, on `ratio == 4` layers;
14//! * the fp32 **compress-state rings** -- `ring_size` slots per *window*
15//!   page, one for the attention compressor and (on ratio-4 layers) one
16//!   for the indexer's own compressor.
17//!
18//! # The window tier is sized independently, and the rest is not
19//!
20//! The window tier reads only the last `P` positions, so it needs
21//! `swa_ratio` of the history and no more. The compressed and indexer
22//! tiers answer questions about the *whole* history, so they stay
23//! anchored to it: `cmp_blocks = full_token / ratio`, with `full_token`
24//! the budget anchor and `swa_ratio` nowhere in it.
25//!
26//! That split is why [`dsv4_pool_sizes`] takes the window size as an
27//! explicit page count instead of deriving everything from one number,
28//! and why the working-set floor is applied **once, in pages**, by the
29//! caller. Applying it by raising `swa_ratio` until the window clears the
30//! floor looks equivalent at one anchor and is not: the ratio keeps
31//! scaling, so at every larger anchor the window is bigger than the floor
32//! ever asked for, and the pool that was sized is not the pool that fits.
33//!
34//! # Why the page solve is a binary search
35//!
36//! [`dsv4_cache_per_page`] collapses all four tiers into one
37//! bytes-per-`P`-tokens number, which is the right shape for a planner
38//! but the wrong tool for the final answer. Dividing the budget by it
39//! assumes every tier scales with the anchor -- and at a small budget the
40//! window does not: it pins at its floor while the full anchor shrinks
41//! underneath it. So [`dsv4_solve_num_pages`] binary-searches the exact
42//! [`dsv4_pool_bytes`] instead. A division there over-sizes the anchor,
43//! every tier inflates past the budget at once, and the failure lands at
44//! the first allocation rather than at config time.
45//!
46//! # Why the window allocator is page-atomic
47//!
48//! The compress-state ring is addressed by
49//! `state_loc = (ws / P) * ring_size + ws % ring_size`, **derived** from
50//! the window slot and never stored. `ring_size` divides `P`, so distinct
51//! window *pages* land on disjoint ring blocks -- but only while every
52//! window slot handed out sits at a page base plus its in-page offset.
53//! [`FreeListAllocator`] therefore hands out unit bases that are always
54//! multiples of the page unit, and [`Dsv4WindowPool::alloc_swa`]
55//! preserves in-page offsets. Break either and two full pages share a
56//! ring block: one request reads another's carry state, with no error
57//! anywhere.
58//!
59//! [`crate::window_pool`] is the generic per-token FIFO window pool --
60//! correct for the pool it ports, and unable to serve a per-page ring,
61//! because a slot it hands back is any free slot rather than a page base.
62//!
63//! # The conservation invariant
64//!
65//! `free units + bound pages == capacity units`, checked by
66//! [`Dsv4WindowPool::check_integrity`]. An **equality**, not a bound: a
67//! `<=` tolerates a leaked window page, which surfaces an hour later as a
68//! pool that cannot admit anything, with nothing to point at.
69//!
70//! Ported 1:1 from FreeToken's `kvcache/dsv4_cost_model.py`,
71//! `kvcache/dsv4_paged_pool.py` and the ring context in
72//! `attention/dsv4_sparse.py` (Apache-2.0); see
73//! `docs/THIRD_PARTY_NOTICES.md`.
74
75use std::collections::BTreeMap;
76
77/// KV, compressed KV and indexer keys are all `bf16`. The fp4/fp8 quant
78/// is an in-place round-trip already baked into the `bf16` value, so
79/// there is no narrower width to price here.
80pub const BF16_BYTES: u64 = 2;
81
82/// The compress-state rings are fp32 -- and only they are.
83pub const FP32_BYTES: u64 = 4;
84
85/// One `int64` per full-history token in the full -> window mapping.
86pub const INT64_BYTES: u64 = 8;
87
88/// `P`: the window page. It is the sliding window *and* the radix block
89/// key, deliberately independent of the generic `page_size`.
90pub const DEFAULT_WINDOW_PAGE: usize = 128;
91
92/// Slack the auto planner reserves on top of the working set, absorbing
93/// plan-vs-measured drift (observed ~265 MiB) and leaving a usable pool.
94pub const AUTO_KV_SLACK_BYTES: u64 = 2 << 30;
95
96/// The "no live window state" sentinel, in both the full -> window
97/// mapping and every ring location derived from it.
98pub const NO_WINDOW_SLOT: i64 = -1;
99
100/// The model geometry the four tiers are priced from.
101///
102/// `compress_ratios` is per layer: `0` for a layer with no compressed
103/// tier, `4` for a compressed + indexer layer, `128` for a
104/// compressed-only layer. A checkpoint may ship more ratios than the
105/// model has layers (44 for 43), so only the first `n_layers` are ever
106/// read -- see [`Dsv4Args::ratios`].
107#[derive(Debug, Clone, PartialEq, Eq, Default)]
108pub struct Dsv4Args {
109    /// The MLA latent width: one slab per token, because V aliases K.
110    pub head_dim: u64,
111    /// The Lightning-Indexer key width.
112    pub index_head_dim: u64,
113    pub n_layers: usize,
114    pub compress_ratios: Vec<u32>,
115}
116
117impl Dsv4Args {
118    /// The ratios the model actually uses: the first `n_layers`.
119    ///
120    /// Truncating rather than cycling or padding, exactly as upstream:
121    /// a checkpoint that ships fewer ratios than layers describes fewer
122    /// layers, and inventing ratios for the rest would price tiers that
123    /// are never allocated.
124    pub fn ratios(&self) -> &[u32] {
125        let end = self.n_layers.min(self.compress_ratios.len());
126        &self.compress_ratios[..end]
127    }
128}
129
130/// Compress-state ring slots per window page (non-speculative).
131///
132/// # Panics
133///
134/// On any ratio other than 4 or 128. The ring geometry is not derivable
135/// from the ratio -- it is a fixed property of the two compressors the
136/// stack ships -- so an unknown ratio is a configuration bug, and
137/// guessing a ring size for it would silently mis-address every carry
138/// state on that layer.
139pub fn ring_size_for_ratio(ratio: u32) -> usize {
140    match ratio {
141        4 => 8,
142        128 => 128,
143        _ => panic!("no ring for ratio {ratio} (only 4 / 128)"),
144    }
145}
146
147/// CSA's ratio. Overlapping blocks, a doubled compressor projection,
148/// and a Lightning Indexer.
149pub const CSA_RATIO: u32 = 4;
150
151/// HCA's ratio. Non-overlapping blocks, a single-width compressor, and
152/// no indexer.
153pub const HCA_RATIO: u32 = 128;
154
155/// Which compressor one layer runs, and everything that follows from it.
156///
157/// The ratio is not a tuning knob with a smooth range: it selects one of
158/// three *different mechanisms*, and the parameters below are not
159/// interpolations between them. Deriving them here, once, is what stops
160/// a single scalar ratio from building an indexer on an HCA layer or
161/// none on a CSA layer -- both of which run, produce numbers, and are
162/// wrong.
163#[derive(Debug, Clone, Copy, PartialEq, Eq)]
164pub enum LayerCompressor {
165    /// Ratio 0: no compressed tier at all. Attention sees the raw
166    /// sliding window and nothing else.
167    ///
168    /// A real entry in the shipped schedule rather than a disabled
169    /// state -- `(0, 0, 4, 128, 4, 128, 4, 0)` opens with two of them
170    /// and closes with one -- so it must be executable, not skipped.
171    None,
172    /// Ratio 4: Compressed Sparse Attention. Overlapping blocks, a
173    /// compressor projection twice as wide, and the Lightning Indexer
174    /// restricting which compressed entries a query may see.
175    Csa,
176    /// Ratio 128: Heavily Compressed Attention. Non-overlapping blocks,
177    /// a single-width compressor, and dense visibility over every
178    /// compressed entry -- no indexer.
179    Hca,
180}
181
182impl LayerCompressor {
183    /// Reads one layer's ratio.
184    ///
185    /// Returns `None` for a ratio that is not 0, 4 or 128, rather than
186    /// approximating it to the nearest mechanism: there is no nearest
187    /// mechanism, and picking one would give that layer the wrong
188    /// compressor width and the wrong indexer, silently.
189    pub fn from_ratio(ratio: u32) -> Option<Self> {
190        match ratio {
191            0 => Some(LayerCompressor::None),
192            CSA_RATIO => Some(LayerCompressor::Csa),
193            HCA_RATIO => Some(LayerCompressor::Hca),
194            _ => None,
195        }
196    }
197
198    /// The ratio this compressor runs at; `0` for [`None`](Self::None).
199    pub fn ratio(self) -> u32 {
200        match self {
201            LayerCompressor::None => 0,
202            LayerCompressor::Csa => CSA_RATIO,
203            LayerCompressor::Hca => HCA_RATIO,
204        }
205    }
206
207    /// Whether this layer instantiates the Lightning Indexer.
208    ///
209    /// CSA only. On an HCA layer every compressed entry is visible, so
210    /// there is nothing for a top-k selector to select; building one
211    /// there costs its own compressor, its own keys and its own tier of
212    /// device memory to answer a question with a fixed answer.
213    pub fn has_indexer(self) -> bool {
214        matches!(self, LayerCompressor::Csa)
215    }
216
217    /// How many times wider this layer's raw per-token compressor
218    /// projection is than one head.
219    ///
220    /// `2` for CSA, and this is the detail a single scalar ratio gets
221    /// wrong on half the stack. Each raw token is projected *twice*:
222    /// once for its role as the tail of the block ending at it, and
223    /// once as the head of the next, overlapping block -- two different
224    /// learned projections of the same token, not one reused twice
225    /// (llama.cpp `load_arch_tensors`' `coff = ratio == 4 ? 2 : 1`, and
226    /// `build_overlap_compressed_kv_from_state`'s
227    /// `GGML_ASSERT(kv_state->ne[0] == 2*n_embd_head)`).
228    pub fn projection_width_multiple(self) -> usize {
229        match self {
230            LayerCompressor::Csa => 2,
231            LayerCompressor::None | LayerCompressor::Hca => 1,
232        }
233    }
234
235    /// Whether consecutive compression blocks share raw positions.
236    pub fn overlapping(self) -> bool {
237        matches!(self, LayerCompressor::Csa)
238    }
239
240    /// How many compressed entries a query at position `pos` may see.
241    ///
242    /// `(pos + 1) / ratio`: ratio-derived, like everything else here,
243    /// and zero on a layer with no compressor. The `+1` is because
244    /// `pos` is an index and the count of tokens through it is one
245    /// more -- without it the query at the last position of a block
246    /// cannot see the block it just completed, which is off by exactly
247    /// one entry for the whole of the sequence.
248    pub fn visible_compressed(self, pos: usize) -> usize {
249        match self.ratio() {
250            0 => 0,
251            r => (pos + 1) / r as usize,
252        }
253    }
254}
255
256/// The schedule a ratio that is not 0, 4 or 128 could not be read into.
257#[derive(Debug, Clone, Copy, PartialEq, Eq)]
258pub struct UnknownCompressRatio {
259    pub layer: usize,
260    pub ratio: u32,
261}
262
263impl std::fmt::Display for UnknownCompressRatio {
264    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
265        let Self { layer, ratio } = self;
266        write!(
267            f,
268            "layer {layer} has compress ratio {ratio}, which is none of 0 (no compressor), \
269             {CSA_RATIO} (CSA) or {HCA_RATIO} (HCA); these are three different mechanisms, \
270             so there is no nearest one to fall back to"
271        )
272    }
273}
274
275impl std::error::Error for UnknownCompressRatio {}
276
277impl Dsv4Args {
278    /// The compressor each layer runs, in layer order.
279    ///
280    /// Reads the same `compress_ratios` the four KV tiers are priced
281    /// from, so what a layer is *sized* for and what it *executes* can
282    /// never drift apart -- which they would the moment the execution
283    /// side kept a scalar of its own.
284    ///
285    /// Refuses the whole schedule on the first unreadable ratio rather
286    /// than dropping that layer: a stack missing one layer's compressor
287    /// still runs, and answers with the wrong attention on it.
288    pub fn compressors(&self) -> Result<Vec<LayerCompressor>, UnknownCompressRatio> {
289        self.ratios()
290            .iter()
291            .enumerate()
292            .map(|(layer, &ratio)| {
293                LayerCompressor::from_ratio(ratio).ok_or(UnknownCompressRatio { layer, ratio })
294            })
295            .collect()
296    }
297}
298
299/// Window pages the sliding pool must always keep for the concurrent
300/// working set.
301///
302/// Each running request's decode transients (2 per request, plus the
303/// dummy's) and, in radix mode, per concurrent request one locked
304/// live-tail page plus a retained soft-pinned prompt-end window -- two
305/// pages, because the retention gap page-aligns to a whole extra page at
306/// `P == window == 128`, so a distinct follow-up per running request can
307/// re-lock two each. Plus the reserved dummy page itself.
308///
309/// The upstream docstring reads "2 per req + dummy", which is
310/// `2 * mr + 1`; the code is `2 * (mr + 1) + 3 * mr + 1`. The **code** is
311/// ported, because it is the one the engine's window floor and the
312/// manager's prefill chunk budget both reserve against -- they have to
313/// agree with each other, not with the prose.
314pub fn dsv4_reserved_window_pages(max_running_req: usize, radix: bool) -> usize {
315    2 * (max_running_req + 1) + if radix { 3 * max_running_req } else { 0 } + 1
316}
317
318/// The only hard floor on DSV4 KV sizing: the live sliding working set
319/// the window pool must always hold.
320///
321/// One prefill chunk's reach, capped at 8 pages (1024 tokens -- chunked
322/// prefill bounds the rest), plus
323/// [`dsv4_reserved_window_pages`]. Everything above it is purely
324/// memory-derived, and a request longer than the anchor is gated
325/// gracefully by the pool's available size rather than by this.
326///
327/// Below the floor a full batch cannot get its window pages at all and
328/// admission deadlocks -- which is why [`dsv4_solve_num_pages`] refuses
329/// at config time instead of letting the pool boot and die at the first
330/// allocation.
331///
332/// `radix` keys on "not naive": DSV4 config resolution rewrites the cache
333/// type to `swa_radix`, so testing for the literal `radix` was always
334/// false upstream.
335pub fn dsv4_window_floor_pages(
336    max_seq_len: usize,
337    max_running_req: usize,
338    radix: bool,
339    page: usize,
340) -> usize {
341    assert!(page > 0, "a window page holds at least one position");
342    let prefill_reach_pages = max_seq_len.div_ceil(page);
343    prefill_reach_pages.min(8) + dsv4_reserved_window_pages(max_running_req, radix)
344}
345
346fn kv_bytes(args: &Dsv4Args) -> u64 {
347    args.head_dim * BF16_BYTES
348}
349
350fn index_bytes(args: &Dsv4Args) -> u64 {
351    args.index_head_dim * BF16_BYTES
352}
353
354/// One attention compress-state row: `kv | score`, each
355/// `(1 + overlap) * head_dim` wide, fp32. Ratio-4 layers overlap.
356fn state_bytes(args: &Dsv4Args, ratio: u32) -> u64 {
357    let overlap = u64::from(ratio == 4);
358    2 * (1 + overlap) * args.head_dim * FP32_BYTES
359}
360
361/// One indexer compress-state row: the same ring geometry keyed on
362/// `index_head_dim`, with overlap always on. Its own pool -- it never
363/// shares slots with the attention ring.
364fn idx_state_bytes(args: &Dsv4Args) -> u64 {
365    2 * 2 * args.index_head_dim * FP32_BYTES
366}
367
368/// `round(ratio * count)` as a count, clamped at zero.
369fn scaled(ratio: f64, count: usize) -> usize {
370    assert!(
371        ratio.is_finite() && ratio >= 0.0,
372        "swa_ratio must be a non-negative fraction of the full history, got {ratio}"
373    );
374    ferrox_core::placement::round_half_even(ratio * count as f64).max(0) as usize
375}
376
377/// Bytes per `P`-token page across **all** tiers, summed over the layers.
378///
379/// This is the one number a budget *division* would use, and it exists
380/// for the affine planner ([`dsv4_auto_cost_model`]) rather than for the
381/// final sizing -- see the module docs on why the solve does not divide.
382///
383/// The window term is scaled by `swa_ratio` and exists on every layer
384/// (all-sliding); the compressed tier is `P / ratio` blocks; the indexer
385/// tier `P / 4` blocks on ratio-4 layers. The two state rings are scaled
386/// by `swa_ratio` **as well**, because they are sized off the *window*
387/// pages (`state_slots = n_win_pages * ring_size`) and not off the full
388/// pages -- charging `ring_size` slots to every full page prices a ring
389/// nobody allocates.
390pub fn dsv4_cache_per_page(args: &Dsv4Args, swa_ratio: f64, page: usize) -> u64 {
391    assert!(page > 0, "a window page holds at least one position");
392    let kv_b = kv_bytes(args);
393    let idx_b = index_bytes(args);
394
395    let mut total = 0u64;
396    for ratio in args.ratios().iter().copied() {
397        // The window tier exists on EVERY layer.
398        total += scaled(swa_ratio, page) as u64 * kv_b;
399        if ratio == 0 {
400            continue;
401        }
402        total += (page as u64 / u64::from(ratio)) * kv_b;
403        if ratio == 4 {
404            total += (page as u64 / 4) * idx_b;
405            total += scaled(swa_ratio, ring_size_for_ratio(4)) as u64 * idx_state_bytes(args);
406        }
407        total += scaled(swa_ratio, ring_size_for_ratio(ratio)) as u64 * state_bytes(args, ratio);
408    }
409    total
410}
411
412/// FULL-tier bytes per full-history token: compressed KV, indexer KV and
413/// the full -> window mapping.
414///
415/// Independent of `swa_ratio` on purpose -- these tiers scale with the
416/// anchor (`cmp_blocks = full_token / ratio`), so their per-token cost
417/// does not move when the window does. The window pool and its rings are
418/// **not** here; see [`dsv4_window_unit_bytes`].
419///
420/// Rounded **up** to whole bytes per token: this is the divisor a cache
421/// slider hands a user, and a conservative maximum is the only kind that
422/// cannot promise a pool the budget will not buy.
423pub fn dsv4_kv_unit_bytes(args: &Dsv4Args, page: usize) -> u64 {
424    assert!(page > 0, "a window page holds at least one position");
425    let kv_b = kv_bytes(args);
426    let idx_b = index_bytes(args);
427    // The full_to_window map: one int64 slot per full token.
428    let mut per_page = page as u64 * INT64_BYTES;
429    for ratio in args.ratios().iter().copied() {
430        if ratio == 0 {
431            continue;
432        }
433        per_page += (page as u64 / u64::from(ratio)) * kv_b;
434        if ratio == 4 {
435            per_page += (page as u64 / 4) * idx_b;
436        }
437    }
438    per_page.div_ceil(page as u64)
439}
440
441/// WINDOW-tier bytes per *window* token: the sliding KV pool on every
442/// layer, plus both state rings (sized off the window pages).
443///
444/// Also independent of `swa_ratio`: the ratio sets how many window tokens
445/// exist, never what one costs.
446pub fn dsv4_window_unit_bytes(args: &Dsv4Args, page: usize) -> u64 {
447    assert!(page > 0, "a window page holds at least one position");
448    let kv_b = kv_bytes(args);
449    let ratios = args.ratios();
450    // Window KV: P slots per page, every layer.
451    let mut per_page = ratios.len() as u64 * page as u64 * kv_b;
452    for ratio in ratios.iter().copied() {
453        if ratio == 0 {
454            continue;
455        }
456        per_page += ring_size_for_ratio(ratio) as u64 * state_bytes(args, ratio);
457        if ratio == 4 {
458            per_page += ring_size_for_ratio(4) as u64 * idx_state_bytes(args);
459        }
460    }
461    per_page.div_ceil(page as u64)
462}
463
464/// One layer's tier sizes. Absent (`None` on
465/// [`Dsv4PoolSizes::layers`]) for a `ratio == 0` layer, which has a
466/// window tier and nothing else.
467#[derive(Debug, Clone, Copy, PartialEq, Eq)]
468pub struct Dsv4LayerSizes {
469    pub ratio: u32,
470    /// Attention compress-state ring slots per window page.
471    pub ring_size: usize,
472    /// Compressed KV blocks: anchored to the FULL history.
473    pub cmp_blocks: usize,
474    /// Indexer KV blocks, on ratio-4 layers only. Also full-anchored.
475    pub idx_blocks: Option<usize>,
476    /// Attention ring slots: `n_win_pages * ring_size` -- window-anchored.
477    pub state_slots: usize,
478    /// Indexer ring slots, on ratio-4 layers only. Window-anchored.
479    pub idx_state_slots: Option<usize>,
480}
481
482/// Per-tier slot counts derived from the budget anchor `full_token`.
483#[derive(Debug, Clone, PartialEq)]
484pub struct Dsv4PoolSizes {
485    /// `P`, the window page.
486    pub page: usize,
487    pub swa_ratio: f64,
488    /// `num_pages * P` -- the budget anchor.
489    pub full_token: usize,
490    /// Global window pool rows.
491    pub n_win_slots: usize,
492    /// `n_win_slots / P`.
493    pub n_win_pages: usize,
494    /// One entry per layer, in layer order.
495    pub layers: Vec<Option<Dsv4LayerSizes>>,
496}
497
498impl Dsv4PoolSizes {
499    /// The anchor in pages: `full_token / P`.
500    pub fn num_pages(&self) -> usize {
501        self.full_token / self.page
502    }
503}
504
505/// Size every tier from `num_pages`, with the window tier sized
506/// independently.
507///
508/// `n_win_pages` is the window in whole pages. `None` means "take
509/// `swa_ratio` of the full history, rounded **up** to whole pages";
510/// `Some` is the caller's own count -- the working-set floor, a pinned
511/// window, or a rebuild's target. Either way it is capped at `num_pages`,
512/// because a window longer than the history it slides over is bytes
513/// nobody can address.
514///
515/// The floor belongs in that `Some`, applied exactly once, in pages. It
516/// must never be applied by raising `swa_ratio`: the ratio multiplies the
517/// anchor, so a ratio raised until it clears the floor at one anchor
518/// overshoots it at every larger one, and the compressed and indexer
519/// tiers -- which stay anchored to the full history here -- then get sized
520/// against a window that was never budgeted.
521///
522/// # Panics
523///
524/// When a non-zero ratio does not divide `page`: `P / ratio` blocks per
525/// page is the compressed tier's whole addressing scheme, and a ragged
526/// division rounds some layer's blocks to zero.
527pub fn dsv4_pool_sizes(
528    num_pages: usize,
529    args: &Dsv4Args,
530    swa_ratio: f64,
531    page: usize,
532    n_win_pages: Option<usize>,
533) -> Dsv4PoolSizes {
534    assert!(page > 0, "a window page holds at least one position");
535    let full_token = num_pages * page;
536
537    let n_win_pages = match n_win_pages {
538        Some(pages) => pages,
539        None => scaled(swa_ratio, full_token).div_ceil(page),
540    };
541    let n_win_pages = n_win_pages.min(num_pages);
542    let n_win_slots = n_win_pages * page;
543
544    let mut layers = Vec::with_capacity(args.ratios().len());
545    for ratio in args.ratios().iter().copied() {
546        if ratio == 0 {
547            layers.push(None);
548            continue;
549        }
550        assert!(
551            page.is_multiple_of(ratio as usize),
552            "P={page} must be divisible by ratio {ratio}"
553        );
554        let ring_size = ring_size_for_ratio(ratio);
555        layers.push(Some(Dsv4LayerSizes {
556            ratio,
557            ring_size,
558            // Full-anchored: the compressed tier answers for the whole
559            // history, whatever fraction of it the window holds.
560            cmp_blocks: full_token / ratio as usize,
561            idx_blocks: (ratio == 4).then_some(full_token / 4),
562            // Window-anchored: one ring block per window page.
563            state_slots: n_win_pages * ring_size,
564            idx_state_slots: (ratio == 4).then(|| n_win_pages * ring_size_for_ratio(4)),
565        }));
566    }
567
568    Dsv4PoolSizes {
569        page,
570        swa_ratio,
571        full_token,
572        n_win_slots,
573        n_win_pages,
574        layers,
575    }
576}
577
578/// The exact bytes a pool built from `sizes` allocates.
579///
580/// Every scratch and sentinel row is counted, because every one of them
581/// is allocated: `n_scratch` rows per compressed and indexer pool (one
582/// per running request row, so a decode whose block did not complete this
583/// step scatters to its own discarded row instead of colliding), one
584/// scratch row per ring, and one sentinel row on the full -> window map.
585/// Leaving them out prices a pool a few rows smaller than the one that
586/// gets built, which is exactly the shortfall a byte-exact solve exists
587/// to avoid.
588///
589/// # Panics
590///
591/// When `sizes` was not built from `args`: the layer counts must agree,
592/// or the window term is summed over a different stack than the tiers.
593pub fn dsv4_pool_bytes(sizes: &Dsv4PoolSizes, args: &Dsv4Args, n_scratch: usize) -> u64 {
594    let ratios = args.ratios();
595    assert_eq!(
596        sizes.layers.len(),
597        ratios.len(),
598        "these sizes were built for a {}-layer stack, not a {}-layer one",
599        sizes.layers.len(),
600        ratios.len()
601    );
602    let kv_b = kv_bytes(args);
603    let idx_b = index_bytes(args);
604    let n_scratch = n_scratch as u64;
605
606    // The window pool, on every layer.
607    let mut total = ratios.len() as u64 * sizes.n_win_slots as u64 * kv_b;
608    // full_to_window, plus its permanent -1 sentinel row.
609    total += (sizes.full_token as u64 + 1) * INT64_BYTES;
610    for layer in sizes.layers.iter().flatten() {
611        total += (layer.cmp_blocks as u64 + n_scratch) * kv_b;
612        total += (layer.state_slots as u64 + 1) * state_bytes(args, layer.ratio);
613        if layer.ratio == 4 {
614            let idx_blocks = layer
615                .idx_blocks
616                .expect("a ratio-4 layer has an indexer tier");
617            let idx_state = layer
618                .idx_state_slots
619                .expect("a ratio-4 layer has an indexer ring");
620            total += (idx_blocks as u64 + n_scratch) * idx_b;
621            total += (idx_state as u64 + 1) * idx_state_bytes(args);
622        }
623    }
624    total
625}
626
627/// The budget cannot pay for the smallest pool that could serve
628/// anything. Refused at config time, before a pool exists.
629#[derive(Debug, Clone, Copy, PartialEq, Eq)]
630pub struct Dsv4BudgetTooSmall {
631    pub available_bytes: u64,
632    /// What the minimal pool costs.
633    pub needed_bytes: u64,
634    /// The minimal anchor, in `P`-pages.
635    pub min_pages: usize,
636    /// The window working-set floor inside it.
637    pub floor_win_pages: usize,
638}
639
640impl std::fmt::Display for Dsv4BudgetTooSmall {
641    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
642        write!(
643            f,
644            "DSV4 KV budget {} bytes cannot fit the minimal pool ({} pages incl. the window \
645             working-set floor {}, needing {} bytes); raise memory_ratio or lower \
646             max_running_req/max_seq_len",
647            self.available_bytes, self.min_pages, self.floor_win_pages, self.needed_bytes
648        )
649    }
650}
651
652impl std::error::Error for Dsv4BudgetTooSmall {}
653
654/// The largest budget-respecting pool: the greatest `num_pages` whose
655/// exact [`dsv4_pool_bytes`] still fits, with the window at
656/// `max(floor_win_pages, ceil(swa_ratio * num_pages))`.
657///
658/// A **binary search over the exact bytes**, not a division. The two are
659/// not the same function: `available / cache_per_page` assumes every tier
660/// scales with the anchor, and at a small budget the window does not --
661/// it pins at its floor while the full anchor shrinks underneath it. The
662/// division therefore returns an anchor at which every tier is inflated
663/// past the budget at once, and the first allocation, not this call, is
664/// where that is discovered.
665///
666/// The floor is honoured in **pages**, once, per
667/// [`dsv4_pool_sizes`]. Below it a full batch cannot get its window pages
668/// and admission deadlocks, so a budget that cannot reach it is an error
669/// here rather than a pool that boots and dies.
670pub fn dsv4_solve_num_pages(
671    available_bytes: u64,
672    args: &Dsv4Args,
673    swa_ratio: f64,
674    floor_win_pages: usize,
675    page: usize,
676    n_scratch: usize,
677) -> Result<Dsv4PoolSizes, Dsv4BudgetTooSmall> {
678    let sizes_at = |num: usize| -> Dsv4PoolSizes {
679        let win = floor_win_pages.max(scaled(swa_ratio, num * page).div_ceil(page));
680        dsv4_pool_sizes(num, args, swa_ratio, page, Some(win))
681    };
682
683    // The full history must at least cover the window working set.
684    let lo0 = floor_win_pages.max(2);
685    let needed = dsv4_pool_bytes(&sizes_at(lo0), args, n_scratch);
686    if needed > available_bytes {
687        return Err(Dsv4BudgetTooSmall {
688            available_bytes,
689            needed_bytes: needed,
690            min_pages: lo0,
691            floor_win_pages,
692        });
693    }
694
695    let mut lo = lo0;
696    // A cheap upper bracket: cache_per_page at ratio 0 undercounts the
697    // window term, so the quotient over-estimates -- then double until it
698    // genuinely does not fit.
699    let mut hi = lo.max((available_bytes / dsv4_cache_per_page(args, 0.0, page).max(1)) as usize);
700    while dsv4_pool_bytes(&sizes_at(hi), args, n_scratch) <= available_bytes {
701        hi *= 2;
702    }
703    while lo < hi - 1 {
704        let mid = (lo + hi) / 2;
705        if dsv4_pool_bytes(&sizes_at(mid), args, n_scratch) <= available_bytes {
706            lo = mid;
707        } else {
708            hi = mid;
709        }
710    }
711    Ok(sizes_at(lo))
712}
713
714/// The affine price of a DSV4 geometry for the MoE-first auto planner.
715#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
716pub struct Dsv4AutoCost {
717    /// Exact marginal cost of one more `P`-page, across all tiers plus
718    /// the full -> window mapping.
719    pub cache_per_page: u64,
720    /// The intercept, anchored at the minimal viable pool.
721    pub fixed_cache_size: u64,
722    /// The reserve floor: the window working set plus slack for
723    /// plan-vs-measured drift, in tokens.
724    pub min_reserve_tokens: usize,
725}
726
727/// The `(per page, fixed, reserve)` triple the MoE-first planner splits
728/// VRAM with.
729///
730/// Conservative at the shipped `swa_ratio`; an extreme ratio can dip
731/// under 1% below exact, which is harmless because `num_pages` is
732/// re-solved byte-exactly by [`dsv4_solve_num_pages`] from the *measured*
733/// memory afterwards. This is the estimate the split is planned with, not
734/// the number the pool is built from.
735pub fn dsv4_auto_cost_model(
736    args: &Dsv4Args,
737    swa_ratio: f64,
738    floor_win_pages: usize,
739    page: usize,
740    n_scratch: usize,
741) -> Dsv4AutoCost {
742    let per_page = dsv4_cache_per_page(args, swa_ratio, page) + page as u64 * INT64_BYTES;
743    let n0 = floor_win_pages.max(2);
744    let win0 = floor_win_pages.max(scaled(swa_ratio, n0 * page).div_ceil(page));
745    let base = dsv4_pool_bytes(
746        &dsv4_pool_sizes(n0, args, swa_ratio, page, Some(win0)),
747        args,
748        n_scratch,
749    );
750    let slack_pages = AUTO_KV_SLACK_BYTES.div_ceil(per_page.max(1)) as usize;
751    Dsv4AutoCost {
752        cache_per_page: per_page,
753        // Saturating: an intercept below zero is a geometry whose marginal
754        // page costs more than the minimal pool does, and a wrapped-around
755        // enormous fixed term would refuse every split.
756        fixed_cache_size: base.saturating_sub(n0 as u64 * per_page),
757        min_reserve_tokens: (n0 + slack_pages) * page,
758    }
759}
760
761// ---------------------------------------------------------------------
762// The paged window allocator
763// ---------------------------------------------------------------------
764
765/// The free list had fewer units than the request needed. Nothing was
766/// taken: the check happens before the first unit moves, so a refused
767/// allocation leaves the allocator exactly as it was and the caller can
768/// evict and retry.
769#[derive(Debug, Clone, Copy, PartialEq, Eq)]
770pub struct FreeListExhausted {
771    pub needed_units: usize,
772    pub available_units: usize,
773    pub capacity: usize,
774    pub page_unit: usize,
775}
776
777impl std::fmt::Display for FreeListExhausted {
778    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
779        write!(
780            f,
781            "window free list out of slots: requested {} units, have {} (capacity {}, unit {})",
782            self.needed_units, self.available_units, self.capacity, self.page_unit
783        )
784    }
785}
786
787impl std::error::Error for FreeListExhausted {}
788
789/// A LIFO free list over **unit bases**, each a multiple of `page_unit`.
790///
791/// The window tier allocates one unit of `page_unit == P` slots at a
792/// time, so every base it hands out is a page base. That is not a
793/// convenience: the compress-state ring block a window slot maps to is
794/// `(ws / P) * ring_size`, so a base that is not a multiple of `P` puts
795/// two full pages' states in one ring block, where each silently
796/// overwrites the other's carry. Keeping the free list in units rather
797/// than slots makes that invariant unbreakable rather than merely
798/// checked.
799///
800/// LIFO -- pop from the tail -- so a page just freed is the next one
801/// handed out, keeping the live window pages clustered.
802///
803/// The compressed and indexer tiers have no allocator at all: their rows
804/// are `full_loc / ratio`, pure arithmetic.
805#[derive(Debug, Clone)]
806pub struct FreeListAllocator {
807    capacity: usize,
808    page_unit: usize,
809    free: Vec<usize>,
810}
811
812impl FreeListAllocator {
813    /// A free list over `capacity` slots in `page_unit`-slot units.
814    ///
815    /// # Panics
816    ///
817    /// When the capacity is not a whole number of units. A ragged tail
818    /// would either be handed out as a short unit or silently lost.
819    pub fn new(capacity: usize, page_unit: usize) -> Self {
820        assert!(page_unit > 0, "a unit spans at least one slot");
821        assert!(
822            capacity.is_multiple_of(page_unit),
823            "capacity {capacity} must be a multiple of page_unit {page_unit}"
824        );
825        let mut allocator = FreeListAllocator {
826            capacity,
827            page_unit,
828            free: Vec::new(),
829        };
830        allocator.reset();
831        allocator
832    }
833
834    /// `n_units` unit bases, each a multiple of `page_unit`.
835    ///
836    /// All or nothing: the capacity check precedes the first pop.
837    pub fn alloc(&mut self, n_units: usize) -> Result<Vec<usize>, FreeListExhausted> {
838        if n_units > self.free.len() {
839            return Err(FreeListExhausted {
840                needed_units: n_units,
841                available_units: self.free.len(),
842                capacity: self.capacity,
843                page_unit: self.page_unit,
844            });
845        }
846        // Take the tail, in ascending order, exactly as the reference
847        // slices it: the caller pairs unit `i` with its `i`-th page.
848        Ok(self.free.split_off(self.free.len() - n_units))
849    }
850
851    /// Return previously-allocated unit bases.
852    ///
853    /// # Panics
854    ///
855    /// On anything that is not a unit base inside the capacity. A base
856    /// that is not page-aligned poisons every later allocation -- and the
857    /// ring aliasing it causes has no symptom other than wrong logits.
858    /// Asserted here rather than assumed, because the callers derive
859    /// these bases from a mapping they may have mutated.
860    pub fn free(&mut self, units: &[usize]) {
861        for base in units.iter().copied() {
862            assert!(
863                base.is_multiple_of(self.page_unit) && base < self.capacity,
864                "{base} is not a unit base of a {}-slot unit inside a capacity of {}",
865                self.page_unit,
866                self.capacity
867            );
868            self.free.push(base);
869        }
870    }
871
872    /// Free capacity in **slots**.
873    pub fn available(&self) -> usize {
874        self.free.len() * self.page_unit
875    }
876
877    /// Free capacity in units.
878    pub fn free_units(&self) -> usize {
879        self.free.len()
880    }
881
882    /// Total slots, allocated or not.
883    pub fn capacity(&self) -> usize {
884        self.capacity
885    }
886
887    pub fn page_unit(&self) -> usize {
888        self.page_unit
889    }
890
891    /// Drop every outstanding allocation.
892    pub fn reset(&mut self) {
893        let n_units = self.capacity / self.page_unit;
894        self.free = (0..n_units).map(|unit| unit * self.page_unit).collect();
895    }
896}
897
898/// The layer-invariant ring context for one decode step of one request.
899#[derive(Debug, Clone, PartialEq, Eq)]
900pub struct Dsv4WindowCtx {
901    /// Where this step's own KV is written.
902    pub window_slot: i64,
903    /// The previous position's slot, for the compressor's carry.
904    pub prev_window_slot: i64,
905    /// One entry per ring slot `j`: the window slot holding the latest
906    /// position `p <= pos` with `p % win == j`, or
907    /// [`NO_WINDOW_SLOT`] where the sequence has not reached it.
908    pub window_slots_topk: Vec<i64>,
909}
910
911/// Which position occupies ring slot `j` at decode position `pos`.
912///
913/// `p = pos - ((pos - j) % win)` -- the latest position at or before
914/// `pos` that is congruent to `j`. `None` before the sequence has reached
915/// the slot (`j > pos`, or `p < 0` early in a decode), which the caller
916/// renders as [`NO_WINDOW_SLOT`] and the sparse kernel masks.
917///
918/// The modulo is **euclidean**: Rust's `%` keeps the sign of the
919/// dividend, so for `j > pos` a truncated remainder names a position
920/// *after* `pos` -- a slot the request has not written yet, read as
921/// though it had.
922pub fn window_ring_position(pos: i64, j: usize, win: usize) -> Option<i64> {
923    assert!(win > 0, "a ring holds at least one slot");
924    if (j as i64) > pos {
925        return None;
926    }
927    let p = pos - (pos - j as i64).rem_euclid(win as i64);
928    (p >= 0).then_some(p)
929}
930
931/// The window tier's page-atomic bookkeeping: which window page backs
932/// each full page, and which window pages are still free.
933///
934/// Token-faced on the outside (the generic cache manager speaks tokens),
935/// page-atomic inside. Window pages are bound 1:1 to full pages and the
936/// per-page state ring requires exactly that, so the page-completeness of
937/// every call is asserted here rather than assumed.
938#[derive(Debug, Clone)]
939pub struct Dsv4WindowPool {
940    page: usize,
941    full_token: usize,
942    n_win_slots: usize,
943    /// Full loc -> window slot, `-1` unbound. One row longer than the
944    /// history: the trailing row is a permanent `-1` sentinel, so a
945    /// gather at `-1` reads it and returns `-1` instead of faulting.
946    full_to_window: Vec<i64>,
947    allocator: FreeListAllocator,
948    chunk_budget: usize,
949}
950
951impl Dsv4WindowPool {
952    /// Build the pool-owned free list and the tail dummy binding.
953    ///
954    /// The **last** full page and the **last** window page are the
955    /// reserved dummy region: the page table's dummy row points at
956    /// `full_token - P`, permanently bound, so graph-padded rows scatter
957    /// to a real slot instead of a negative index. That page is bound
958    /// outside the free list -- the allocator's capacity is
959    /// `n_win_slots - P` -- which is also what makes the usable window
960    /// count `swa_num_tokens - 1`, the same capacity convention the
961    /// generic pool gets from reserving slot 0.
962    pub fn new(sizes: &Dsv4PoolSizes, max_running_req: usize, radix: bool) -> Self {
963        let page = sizes.page;
964        assert!(
965            sizes.full_token >= page && sizes.full_token.is_multiple_of(page),
966            "the full anchor must be whole pages and hold the dummy page"
967        );
968        assert!(
969            sizes.n_win_slots >= page && sizes.n_win_slots.is_multiple_of(page),
970            "the window pool must be whole pages and hold the dummy page"
971        );
972        let mut pool = Dsv4WindowPool {
973            page,
974            full_token: sizes.full_token,
975            n_win_slots: sizes.n_win_slots,
976            full_to_window: vec![NO_WINDOW_SLOT; sizes.full_token + 1],
977            allocator: FreeListAllocator::new(sizes.n_win_slots - page, page),
978            chunk_budget: 0,
979        };
980        pool.bind_window_pages(sizes.full_token - page, sizes.n_win_slots - page);
981
982        // A batched prefill holds the whole chunk's window live at once
983        // (sliding frees only between chunks, so the peak is ~2x the
984        // chunk): reserve the concurrent working set and halve the rest.
985        let n_win_pages = (sizes.n_win_slots / page) - 1;
986        let reserved = dsv4_reserved_window_pages(max_running_req, radix);
987        pool.chunk_budget = page.max(n_win_pages.saturating_sub(reserved) / 2 * page);
988        pool
989    }
990
991    /// `P`: the window page, the sliding window, and the radix block key.
992    pub fn page_size(&self) -> usize {
993        self.page
994    }
995
996    /// The window pool in the generic capacity convention: allocatable
997    /// slots plus one. The dummy page is already excluded from the free
998    /// list, so the `+1` re-encodes the same `capacity == tokens - 1`
999    /// the generic pool gets from its slot-0 sentinel.
1000    pub fn swa_num_tokens(&self) -> usize {
1001        (self.n_win_slots - self.page) + 1
1002    }
1003
1004    /// Window slots that can still be handed out.
1005    pub fn swa_available_size(&self) -> usize {
1006        self.allocator.available()
1007    }
1008
1009    /// The largest prefill chunk this pool can hold the window for.
1010    pub fn prefill_chunk_budget(&self) -> usize {
1011        self.chunk_budget
1012    }
1013
1014    /// Permanently bind one full page to one window page, offsets
1015    /// preserved.
1016    ///
1017    /// # Panics
1018    ///
1019    /// On a base that is not page-aligned: the ring block layout is keyed
1020    /// on the page base, so an unaligned binding aliases two pages onto
1021    /// one ring block.
1022    pub fn bind_window_pages(&mut self, full_page_base: usize, window_page_base: usize) {
1023        assert!(
1024            full_page_base.is_multiple_of(self.page) && window_page_base.is_multiple_of(self.page),
1025            "window bindings are page-aligned: full {full_page_base}, window {window_page_base}"
1026        );
1027        for offset in 0..self.page {
1028            self.full_to_window[full_page_base + offset] = (window_page_base + offset) as i64;
1029        }
1030    }
1031
1032    /// Drop the binding of these full locs, returning nothing to the
1033    /// free list. Negative locs are ignored.
1034    pub fn unbind_window_pages(&mut self, full_locs: &[i64]) {
1035        for loc in full_locs.iter().copied().filter(|loc| *loc >= 0) {
1036            self.full_to_window[loc as usize] = NO_WINDOW_SLOT;
1037        }
1038    }
1039
1040    /// Bind one window page per incoming full page.
1041    ///
1042    /// `full_indices` must be whole contiguous ascending pages -- the
1043    /// page-to-token expansion the caller already performs. The in-page
1044    /// offsets are **preserved** (`window_slot = wbase + pos % P`), which
1045    /// is what the state ring's page-block layout requires: a slot
1046    /// permuted inside its page still lands in the right ring block, but
1047    /// on the wrong slot of it.
1048    ///
1049    /// All or nothing: the pages are allocated before the first mapping
1050    /// row is written, so an exhausted pool leaves the mapping untouched.
1051    ///
1052    /// # Panics
1053    ///
1054    /// On a partial, unaligned or non-ascending page. Upstream gets these
1055    /// from the page expansion by construction; asserted here, not
1056    /// assumed.
1057    pub fn alloc_swa(&mut self, full_indices: &[i64]) -> Result<(), FreeListExhausted> {
1058        if full_indices.is_empty() {
1059            return Ok(());
1060        }
1061        let page = self.page;
1062        assert!(
1063            full_indices.len().is_multiple_of(page),
1064            "alloc_swa needs whole pages, got {} slots",
1065            full_indices.len()
1066        );
1067        let mut bases = Vec::with_capacity(full_indices.len() / page);
1068        for chunk in full_indices.chunks(page) {
1069            let base = chunk[0];
1070            assert!(
1071                base >= 0 && (base as usize).is_multiple_of(page),
1072                "alloc_swa pages start at a page base, got {base}"
1073            );
1074            for (offset, loc) in chunk.iter().copied().enumerate() {
1075                assert_eq!(
1076                    loc,
1077                    base + offset as i64,
1078                    "alloc_swa pages must be contiguous ascending"
1079                );
1080            }
1081            debug_assert_eq!(
1082                self.full_to_window[base as usize], NO_WINDOW_SLOT,
1083                "full page {base} already holds a window page; binding over it would leak it"
1084            );
1085            bases.push(base as usize);
1086        }
1087
1088        let wbases = self.allocator.alloc(bases.len())?;
1089        for (fbase, wbase) in bases.into_iter().zip(wbases) {
1090            for offset in 0..page {
1091                self.full_to_window[fbase + offset] = (wbase + offset) as i64;
1092            }
1093        }
1094        Ok(())
1095    }
1096
1097    /// Return the window pages backing these full locs and unbind them.
1098    ///
1099    /// Page-atomic: the locs must cover every touched page completely.
1100    /// Idempotent over pages that are already unbound -- a slide, a
1101    /// tombstone and an eviction pass may all name the same page, and
1102    /// only the first hands its window page back.
1103    ///
1104    /// # Panics
1105    ///
1106    /// On a partially covered page. Freeing half a page would return a
1107    /// window page whose other half is still mapped, so the next
1108    /// allocation gets a page two full pages read through.
1109    pub fn free_swa(&mut self, full_indices: &[i64]) {
1110        let page = self.page;
1111        let mut counts: BTreeMap<usize, usize> = BTreeMap::new();
1112        for loc in full_indices.iter().copied().filter(|loc| *loc >= 0) {
1113            *counts.entry(loc as usize / page * page).or_insert(0) += 1;
1114        }
1115        if counts.is_empty() {
1116            return;
1117        }
1118        let partial: Vec<(usize, usize)> = counts
1119            .iter()
1120            .filter(|(_, count)| **count != page)
1121            .map(|(base, count)| (*base, *count))
1122            .take(4)
1123            .collect();
1124        assert!(
1125            partial.is_empty(),
1126            "free_swa got partial pages (base, count): {partial:?}"
1127        );
1128
1129        let mut freed = Vec::with_capacity(counts.len());
1130        for base in counts.into_keys() {
1131            let window_slot = self.full_to_window[base];
1132            for offset in 0..page {
1133                self.full_to_window[base + offset] = NO_WINDOW_SLOT;
1134            }
1135            if window_slot >= 0 {
1136                freed.push(window_slot as usize / page * page);
1137            }
1138        }
1139        self.allocator.free(&freed);
1140    }
1141
1142    /// The window slot holding `full_loc`, or [`NO_WINDOW_SLOT`].
1143    ///
1144    /// A negative loc reads the permanent sentinel row and returns `-1`,
1145    /// so callers carrying "no such position" as `-1` need no special
1146    /// case -- upstream gets the same effect from negative tensor
1147    /// indexing into the trailing row.
1148    pub fn translate(&self, full_loc: i64) -> i64 {
1149        if full_loc < 0 {
1150            return NO_WINDOW_SLOT;
1151        }
1152        self.full_to_window[full_loc as usize]
1153    }
1154
1155    /// The compress-state ring location for a window slot -- **derived**,
1156    /// never stored.
1157    ///
1158    /// `(ws / P) * ring_size + ws % ring_size`: the page picks the ring
1159    /// block, the in-page offset picks the slot inside it. `ring_size`
1160    /// divides `P`, so distinct window pages land on disjoint blocks.
1161    ///
1162    /// Deriving it on every use is the whole safety property. A stored
1163    /// `state_loc` outlives the binding it was derived from: free a
1164    /// window page and the next request to take it inherits the same ring
1165    /// block, so a cached location reads another request's carry state --
1166    /// no error, no fault, just wrong numbers.
1167    ///
1168    /// # Panics
1169    ///
1170    /// When `ring_size` does not divide `page`, which is what makes the
1171    /// blocks disjoint in the first place.
1172    pub fn state_loc(window_slot: i64, ring_size: usize, page: usize) -> i64 {
1173        assert!(
1174            ring_size > 0 && page.is_multiple_of(ring_size),
1175            "ring_size {ring_size} must divide P={page}, or two pages share a ring block"
1176        );
1177        if window_slot < 0 {
1178            return NO_WINDOW_SLOT;
1179        }
1180        let pages = window_slot / page as i64;
1181        pages * ring_size as i64 + window_slot % ring_size as i64
1182    }
1183
1184    /// The ring context for one decode step: this position's slot, the
1185    /// previous position's, and the whole `P`-slot ring.
1186    ///
1187    /// `full_locs` is the request's snapshot -- full loc per position --
1188    /// read instead of the live mapping so a concurrent allocation cannot
1189    /// redirect an in-flight step. Computed fresh every call: caching it
1190    /// freezes a replay at the capture-time ring slots, `-1` fills
1191    /// included.
1192    pub fn window_ctx(&self, pos: usize, full_locs: &[i64]) -> Dsv4WindowCtx {
1193        let win = self.page;
1194        let window_slots_topk = (0..win)
1195            .map(|j| match window_ring_position(pos as i64, j, win) {
1196                Some(p) => self.translate(full_locs[p as usize]),
1197                None => NO_WINDOW_SLOT,
1198            })
1199            .collect();
1200        Dsv4WindowCtx {
1201            window_slot: self.translate(full_locs[pos]),
1202            prev_window_slot: self.translate(full_locs[pos.saturating_sub(1)]),
1203            window_slots_topk,
1204        }
1205    }
1206
1207    /// Every window page is either free or bound to exactly one full
1208    /// page, and every binding is page-atomic.
1209    ///
1210    /// The count is an **equality**: free units plus bound pages is the
1211    /// capacity exactly. A `<=` would catch a double free and tolerate a
1212    /// leak, and a leaked window page is the failure that shows up an
1213    /// hour later as a pool that admits nothing.
1214    pub fn check_integrity(&self) {
1215        let page = self.page;
1216        let n_full_pages = self.full_token / page;
1217        let mut seen = vec![false; self.n_win_slots / page];
1218        let mut bound = 0usize;
1219
1220        // The dummy page (the last one) is permanently bound outside the
1221        // free list, so it is neither free nor counted.
1222        assert_eq!(
1223            self.full_to_window[self.full_token - page],
1224            (self.n_win_slots - page) as i64,
1225            "the reserved dummy page lost its permanent binding"
1226        );
1227        assert_eq!(
1228            self.full_to_window[self.full_token], NO_WINDOW_SLOT,
1229            "the trailing sentinel row was written"
1230        );
1231
1232        for full_page in 0..n_full_pages - 1 {
1233            let base = full_page * page;
1234            let window_slot = self.full_to_window[base];
1235            for offset in 0..page {
1236                let expected = if window_slot < 0 {
1237                    NO_WINDOW_SLOT
1238                } else {
1239                    window_slot + offset as i64
1240                };
1241                assert_eq!(
1242                    self.full_to_window[base + offset],
1243                    expected,
1244                    "full page {base} is bound partially or out of order at offset {offset}"
1245                );
1246            }
1247            if window_slot < 0 {
1248                continue;
1249            }
1250            assert!(
1251                (window_slot as usize).is_multiple_of(page),
1252                "window base {window_slot} is not page-aligned; its ring block aliases"
1253            );
1254            let index = window_slot as usize / page;
1255            assert!(
1256                !std::mem::replace(&mut seen[index], true),
1257                "window page {window_slot} is bound to two full pages"
1258            );
1259            bound += 1;
1260        }
1261
1262        for base in self.free_bases() {
1263            assert!(
1264                !std::mem::replace(&mut seen[base / page], true),
1265                "window page {base} is both free and bound, or free twice"
1266            );
1267        }
1268
1269        let capacity_units = self.allocator.capacity() / page;
1270        assert_eq!(
1271            self.allocator.free_units() + bound,
1272            capacity_units,
1273            "window pages leaked or double-freed: {} free + {bound} bound != {capacity_units}",
1274            self.allocator.free_units()
1275        );
1276    }
1277
1278    fn free_bases(&self) -> Vec<usize> {
1279        self.allocator.free.clone()
1280    }
1281}
1282
1283#[cfg(test)]
1284mod tests {
1285
1286    /// The shipped schedule, and the whole reason a scalar ratio is
1287    /// wrong: one array holds all three mechanisms, so any single value
1288    /// applied uniformly is right for at most one kind of layer.
1289    #[test]
1290    fn the_shipped_schedule_reads_as_three_different_mechanisms() {
1291        let args = Dsv4Args {
1292            head_dim: 64,
1293            index_head_dim: 32,
1294            n_layers: 8,
1295            compress_ratios: vec![0, 0, 4, 128, 4, 128, 4, 0],
1296        };
1297        assert_eq!(
1298            args.compressors().unwrap(),
1299            vec![
1300                LayerCompressor::None,
1301                LayerCompressor::None,
1302                LayerCompressor::Csa,
1303                LayerCompressor::Hca,
1304                LayerCompressor::Csa,
1305                LayerCompressor::Hca,
1306                LayerCompressor::Csa,
1307                LayerCompressor::None,
1308            ]
1309        );
1310    }
1311
1312    /// The two properties a uniform ratio gets wrong on half the stack:
1313    /// an indexer built where every entry is already visible, and a
1314    /// compressor projection of the wrong width. Both run and produce
1315    /// numbers, which is why they are derived from the mechanism rather
1316    /// than configured.
1317    #[test]
1318    fn the_indexer_and_the_projection_width_follow_from_the_mechanism() {
1319        assert!(LayerCompressor::Csa.has_indexer());
1320        assert!(!LayerCompressor::Hca.has_indexer());
1321        assert!(!LayerCompressor::None.has_indexer());
1322
1323        assert_eq!(LayerCompressor::Csa.projection_width_multiple(), 2);
1324        assert_eq!(LayerCompressor::Hca.projection_width_multiple(), 1);
1325        assert_eq!(LayerCompressor::None.projection_width_multiple(), 1);
1326
1327        assert!(LayerCompressor::Csa.overlapping());
1328        assert!(!LayerCompressor::Hca.overlapping());
1329    }
1330
1331    /// Visibility is ratio-derived, and the `+1` is load-bearing: the
1332    /// query at the last position of a block must see the block it just
1333    /// completed. Without it every layer is short exactly one
1334    /// compressed entry, for the whole sequence.
1335    #[test]
1336    fn a_query_sees_one_compressed_entry_per_completed_block() {
1337        let csa = LayerCompressor::Csa;
1338        assert_eq!(csa.visible_compressed(0), 0, "no block is complete yet");
1339        assert_eq!(csa.visible_compressed(2), 0);
1340        assert_eq!(csa.visible_compressed(3), 1, "the first block just closed");
1341        assert_eq!(csa.visible_compressed(7), 2);
1342        assert_eq!(csa.visible_compressed(8), 2);
1343
1344        let hca = LayerCompressor::Hca;
1345        assert_eq!(hca.visible_compressed(126), 0);
1346        assert_eq!(hca.visible_compressed(127), 1);
1347        assert_eq!(hca.visible_compressed(255), 2);
1348
1349        // A layer with no compressor has nothing to see, at any
1350        // position -- not "all of them", which a ratio of zero would
1351        // produce as a divide by zero rather than an answer.
1352        for pos in [0, 1, 127, 1_000_000] {
1353            assert_eq!(LayerCompressor::None.visible_compressed(pos), 0);
1354        }
1355    }
1356
1357    /// A ratio that is none of the three is refused, naming its layer.
1358    /// There is no nearest mechanism to round to, and picking one gives
1359    /// that layer the wrong compressor width and the wrong indexer with
1360    /// nothing to point at afterwards.
1361    #[test]
1362    fn an_unknown_ratio_is_refused_and_names_its_layer() {
1363        assert_eq!(LayerCompressor::from_ratio(7), None);
1364        assert_eq!(LayerCompressor::from_ratio(64), None);
1365
1366        let args = Dsv4Args {
1367            head_dim: 64,
1368            index_head_dim: 32,
1369            n_layers: 3,
1370            compress_ratios: vec![0, 64, 128],
1371        };
1372        let err = args.compressors().unwrap_err();
1373        assert_eq!(
1374            err,
1375            UnknownCompressRatio {
1376                layer: 1,
1377                ratio: 64
1378            }
1379        );
1380        assert!(err.to_string().contains("layer 1"));
1381    }
1382
1383    /// The schedule reads the same array the tiers are priced from, and
1384    /// truncates the same way -- a checkpoint shipping 44 ratios for 43
1385    /// layers describes 43 layers. Sizing a tier this stack never
1386    /// executes, or executing one it never sized, are the two failures
1387    /// sharing the array prevents.
1388    #[test]
1389    fn the_schedule_and_the_sizing_read_the_same_truncated_array() {
1390        let args = Dsv4Args {
1391            head_dim: 64,
1392            index_head_dim: 32,
1393            n_layers: 3,
1394            compress_ratios: vec![4, 128, 0, 4],
1395        };
1396        assert_eq!(args.ratios(), &[4, 128, 0]);
1397        let compressors = args.compressors().unwrap();
1398        assert_eq!(compressors.len(), args.ratios().len());
1399        for (c, &r) in compressors.iter().zip(args.ratios()) {
1400            assert_eq!(c.ratio(), r);
1401        }
1402    }
1403
1404    /// Every ratio the ring geometry accepts is one the schedule can
1405    /// read, and vice versa for the compressed tiers. The two tables
1406    /// are separate functions over the same three values, and a ratio
1407    /// only one of them knows is a layer that is sized without being
1408    /// executable or the reverse.
1409    #[test]
1410    fn the_ring_table_and_the_compressor_table_agree_on_which_ratios_exist() {
1411        for ratio in [CSA_RATIO, HCA_RATIO] {
1412            assert!(LayerCompressor::from_ratio(ratio).is_some());
1413            assert!(ring_size_for_ratio(ratio) > 0);
1414        }
1415        // Ratio 0 is readable as a mechanism but has no ring, because a
1416        // layer with no compressor has no carry state to address.
1417        assert_eq!(LayerCompressor::from_ratio(0), Some(LayerCompressor::None));
1418        assert!(std::panic::catch_unwind(|| ring_size_for_ratio(0)).is_err());
1419    }
1420    use super::*;
1421
1422    const P: usize = DEFAULT_WINDOW_PAGE;
1423
1424    /// head_dim 8 -> kv 16 B; index_head_dim 4 -> idx 8 B.
1425    /// Ratio-4 state rows 2*2*8*4 = 128 B, ratio-128 rows 2*1*8*4 = 64 B,
1426    /// indexer state rows 2*2*4*4 = 64 B.
1427    fn args() -> Dsv4Args {
1428        Dsv4Args {
1429            head_dim: 8,
1430            index_head_dim: 4,
1431            n_layers: 4,
1432            compress_ratios: vec![0, 4, 128, 4],
1433        }
1434    }
1435
1436    /// `dsv4_pool_bytes` for [`args`], worked out by hand:
1437    /// window 4 layers x 16 B x P = 8192 B per window page, the two
1438    /// ratio-4 layers' rings 2 x (8x128 + 8x64) = 3072 B per window page,
1439    /// the ratio-128 layer's ring 128x64 = 8192 B per window page;
1440    /// per full page 8 x 128 mapping + 2 x (32x16 + 32x8) + 1x16 = 2576 B.
1441    fn expected_bytes(num_pages: usize, win_pages: usize, n_scratch: usize) -> u64 {
1442        19456 * win_pages as u64
1443            + 2576 * num_pages as u64
1444            // scratch rows: 2 x (cmp 16 B + idx 8 B) + 1 x cmp 16 B
1445            + 64 * n_scratch as u64
1446            // ring scratch rows 2x(128+64) + 1x64, mapping sentinel 8
1447            + 456
1448    }
1449
1450    // ---- the tiered cost model ----
1451
1452    #[test]
1453    fn ring_sizes_are_fixed_per_ratio() {
1454        assert_eq!(ring_size_for_ratio(4), 8);
1455        assert_eq!(ring_size_for_ratio(128), 128);
1456    }
1457
1458    /// The ring geometry is a property of the two compressors that ship,
1459    /// not a function of the ratio -- so an unknown ratio is refused
1460    /// rather than given a guessed ring.
1461    #[test]
1462    #[should_panic(expected = "no ring for ratio 8")]
1463    fn an_unsupported_ratio_has_no_ring() {
1464        ring_size_for_ratio(8);
1465    }
1466
1467    /// The docstring says "2 per req + dummy"; the code reserves
1468    /// `2 * (mr + 1) + 3 * mr + 1`. The code is what the engine floor and
1469    /// the chunk budget both reserve against, so the code is what is
1470    /// ported.
1471    #[test]
1472    fn the_reserved_window_pages_follow_the_code_not_the_docstring() {
1473        assert_eq!(dsv4_reserved_window_pages(2, true), 2 * 3 + 3 * 2 + 1);
1474        assert_eq!(dsv4_reserved_window_pages(2, false), 2 * 3 + 1);
1475        // Not the docstring's 2 * mr + 1.
1476        assert_ne!(dsv4_reserved_window_pages(2, false), 2 * 2 + 1);
1477    }
1478
1479    #[test]
1480    fn the_window_floor_caps_the_prefill_reach_at_eight_pages() {
1481        // 2048 tokens is 16 pages of reach, capped at 8.
1482        assert_eq!(
1483            dsv4_window_floor_pages(2048, 2, true, P),
1484            8 + dsv4_reserved_window_pages(2, true)
1485        );
1486        // A short context pays only its own reach.
1487        assert_eq!(
1488            dsv4_window_floor_pages(256, 2, true, P),
1489            2 + dsv4_reserved_window_pages(2, true)
1490        );
1491    }
1492
1493    #[test]
1494    fn the_per_page_cost_sums_every_tier_over_the_layers() {
1495        // window 4 x round(0.5 x 128) x 16 = 4096
1496        // ratio-4 layers x2: 32x16 + 32x8 + round(0.5x8)x64 + round(0.5x8)x128
1497        // ratio-128 layer: 1x16 + round(0.5x128)x64
1498        assert_eq!(
1499            dsv4_cache_per_page(&args(), 0.5, P),
1500            4096 + 2 * (512 + 256 + 256 + 512) + (16 + 4096)
1501        );
1502        // At ratio 0 the window tier and both rings vanish, leaving only
1503        // the full-anchored tiers -- which is the undercounting bracket
1504        // the page solve starts its search from.
1505        assert_eq!(dsv4_cache_per_page(&args(), 0.0, P), 2 * (512 + 256) + 16);
1506    }
1507
1508    /// The sizing rounds halves to even (Python's `round`), while the
1509    /// unit costs round bytes-per-token up. `f64::round` would take
1510    /// 0.5 -> 1 here and buy a ring slot per page the budget never
1511    /// priced.
1512    #[test]
1513    fn the_ratio_scaling_rounds_halves_to_even() {
1514        let args = args();
1515        // The ratio-4 rings: round(0.0625 x 8) == round(0.5) == 0, not 1,
1516        // so those two layers buy no ring slots at all on this page.
1517        let half_down = dsv4_cache_per_page(&args, 0.0625, P);
1518        assert_eq!(half_down, 4 * 8 * 16 + 2 * (512 + 256) + (16 + 8 * 64));
1519        // round(0.1875 x 8) == round(1.5) == 2, up to the even neighbour.
1520        let half_up = dsv4_cache_per_page(&args, 0.1875, P);
1521        assert_eq!(
1522            half_up,
1523            4 * 24 * 16 + 2 * (512 + 256 + 2 * 64 + 2 * 128) + (16 + 24 * 64)
1524        );
1525    }
1526
1527    #[test]
1528    fn the_unit_costs_round_bytes_per_token_up() {
1529        // full tier: 1024 mapping + 2 x 768 + 16 = 2576 per page -> 20.125
1530        assert_eq!(dsv4_kv_unit_bytes(&args(), P), 21);
1531        // window tier: 8192 + 2 x 1536 + 8192 = 19456 per page -> exact
1532        assert_eq!(dsv4_window_unit_bytes(&args(), P), 152);
1533    }
1534
1535    /// The window tier is `swa_ratio` of the history; the compressed and
1536    /// indexer tiers stay anchored to the FULL history whatever the
1537    /// window does.
1538    #[test]
1539    fn the_window_tier_is_sized_independently_of_the_full_anchor() {
1540        let args = args();
1541        let sizes = dsv4_pool_sizes(64, &args, 0.1, P, None);
1542        assert_eq!(sizes.full_token, 64 * P);
1543        // round(0.1 x 8192) = 819 slots -> ceil to 7 whole pages.
1544        assert_eq!(sizes.n_win_pages, 7);
1545        assert_eq!(sizes.n_win_slots, 7 * P);
1546
1547        assert_eq!(sizes.layers[0], None, "a ratio-0 layer has no tiers");
1548        let ratio4 = sizes.layers[1].unwrap();
1549        assert_eq!(ratio4.cmp_blocks, 64 * P / 4, "full-anchored");
1550        assert_eq!(ratio4.idx_blocks, Some(64 * P / 4));
1551        assert_eq!(ratio4.state_slots, 7 * 8, "window-anchored");
1552        assert_eq!(ratio4.idx_state_slots, Some(7 * 8));
1553        let ratio128 = sizes.layers[2].unwrap();
1554        assert_eq!(ratio128.cmp_blocks, 64 * P / 128);
1555        assert_eq!(ratio128.idx_blocks, None);
1556        assert_eq!(ratio128.state_slots, 7 * 128);
1557    }
1558
1559    #[test]
1560    fn an_explicit_window_is_capped_at_the_full_history() {
1561        let sizes = dsv4_pool_sizes(8, &args(), 0.5, P, Some(64));
1562        assert_eq!(
1563            sizes.n_win_pages, 8,
1564            "a window past the history is bytes nobody can address"
1565        );
1566        // And a ratio over 1.0 is capped the same way.
1567        assert_eq!(dsv4_pool_sizes(8, &args(), 4.0, P, None).n_win_pages, 8);
1568    }
1569
1570    #[test]
1571    fn the_pool_bytes_are_the_sum_of_every_allocated_row() {
1572        let args = args();
1573        let sizes = dsv4_pool_sizes(64, &args, 0.1, P, Some(21));
1574        assert_eq!(dsv4_pool_bytes(&sizes, &args, 1), expected_bytes(64, 21, 1));
1575        // Each scratch row is really allocated, so each is really priced.
1576        assert_eq!(
1577            dsv4_pool_bytes(&sizes, &args, 3) - dsv4_pool_bytes(&sizes, &args, 1),
1578            2 * 64
1579        );
1580    }
1581
1582    /// The naive solve -- `available / cache_per_page` -- returns an
1583    /// anchor whose exact bytes are OVER the budget, because the window
1584    /// tier does not scale with the anchor once it pins at its floor.
1585    /// The binary search returns the largest anchor that actually fits.
1586    #[test]
1587    fn dividing_the_budget_by_the_per_page_cost_overshoots_it() {
1588        let args = args();
1589        let floor = 21;
1590        let available = 500_000;
1591        let solved = dsv4_solve_num_pages(available, &args, 0.5, floor, P, 1).expect("fits");
1592
1593        let naive =
1594            (available / (dsv4_cache_per_page(&args, 0.5, P) + P as u64 * INT64_BYTES)) as usize;
1595        assert_eq!(naive, 40);
1596        let naive_sizes = dsv4_pool_sizes(naive, &args, 0.5, P, Some(floor.max(naive.div_ceil(2))));
1597        assert!(
1598            dsv4_pool_bytes(&naive_sizes, &args, 1) > available,
1599            "the division must be the one that overshoots"
1600        );
1601
1602        assert_eq!(solved.num_pages(), 35);
1603        assert!(dsv4_pool_bytes(&solved, &args, 1) <= available);
1604    }
1605
1606    /// At a small budget the window pins at its floor in PAGES and the
1607    /// full anchor shrinks underneath it. Honouring the floor by
1608    /// inflating `swa_ratio` instead -- the ratio that yields the floor at
1609    /// the minimal pool -- keeps the window scaling with the anchor, so
1610    /// at the solved anchor it buys far more window than the floor asked
1611    /// for and blows the budget.
1612    #[test]
1613    fn a_small_budget_pins_the_window_at_its_floor_and_shrinks_the_full_anchor() {
1614        let args = args();
1615        let floor = 21;
1616        let available = 900_000;
1617        let solved = dsv4_solve_num_pages(available, &args, 0.1, floor, P, 1).expect("fits");
1618
1619        assert_eq!(solved.n_win_pages, floor, "the window pinned at its floor");
1620        assert_eq!(solved.num_pages(), 190, "the full anchor took the rest");
1621        assert!(solved.num_pages() > solved.n_win_pages * 4);
1622        // The full-anchored tiers still cover the whole history.
1623        assert_eq!(solved.layers[1].unwrap().cmp_blocks, 190 * P / 4);
1624
1625        let inflated = floor as f64 / floor.max(2) as f64; // 1.0
1626        let inflated_sizes = dsv4_pool_sizes(solved.num_pages(), &args, inflated, P, None);
1627        assert!(inflated_sizes.n_win_pages > solved.n_win_pages);
1628        assert!(
1629            dsv4_pool_bytes(&inflated_sizes, &args, 1) > available,
1630            "inflating swa_ratio to carry the floor busts the budget"
1631        );
1632    }
1633
1634    #[test]
1635    fn the_solved_pool_is_the_largest_that_fits() {
1636        let args = args();
1637        let available = 900_000;
1638        let solved = dsv4_solve_num_pages(available, &args, 0.1, 21, P, 1).expect("fits");
1639        let bytes = dsv4_pool_bytes(&solved, &args, 1);
1640        assert!(bytes <= available);
1641
1642        let one_more = dsv4_pool_sizes(
1643            solved.num_pages() + 1,
1644            &args,
1645            0.1,
1646            P,
1647            Some(21.max(scaled(0.1, (solved.num_pages() + 1) * P).div_ceil(P))),
1648        );
1649        assert!(dsv4_pool_bytes(&one_more, &args, 1) > available);
1650    }
1651
1652    /// Below the floor a full batch cannot get its window pages and
1653    /// admission deadlocks -- so the budget is refused at config time
1654    /// rather than at the first allocation.
1655    #[test]
1656    fn a_budget_below_the_minimal_pool_is_refused_at_config_time() {
1657        let args = args();
1658        let err = dsv4_solve_num_pages(100_000, &args, 0.1, 21, P, 1).unwrap_err();
1659        assert_eq!(err.min_pages, 21);
1660        assert_eq!(err.floor_win_pages, 21);
1661        assert!(err.needed_bytes > err.available_bytes);
1662        assert!(err.to_string().contains("working-set floor 21"));
1663    }
1664
1665    #[test]
1666    fn the_auto_cost_model_is_affine_through_the_minimal_pool() {
1667        let args = args();
1668        let floor = 21;
1669        let cost = dsv4_auto_cost_model(&args, 0.1, floor, P, 1);
1670        assert_eq!(
1671            cost.cache_per_page,
1672            dsv4_cache_per_page(&args, 0.1, P) + P as u64 * INT64_BYTES
1673        );
1674
1675        // The intercept is anchored at the minimal pool: the affine price
1676        // reproduces its exact bytes there, as an equality.
1677        let n0 = floor.max(2);
1678        let base = dsv4_pool_bytes(&dsv4_pool_sizes(n0, &args, 0.1, P, Some(floor)), &args, 1);
1679        assert_eq!(
1680            cost.fixed_cache_size + n0 as u64 * cost.cache_per_page,
1681            base
1682        );
1683
1684        let slack_pages = AUTO_KV_SLACK_BYTES.div_ceil(cost.cache_per_page) as usize;
1685        assert_eq!(cost.min_reserve_tokens, (n0 + slack_pages) * P);
1686    }
1687
1688    // ---- the paged window allocator ----
1689
1690    fn pool() -> Dsv4WindowPool {
1691        // 8 full pages, 5 window pages -> 4 allocatable, 1 dummy.
1692        let sizes = dsv4_pool_sizes(8, &args(), 0.5, P, Some(5));
1693        Dsv4WindowPool::new(&sizes, 2, true)
1694    }
1695
1696    fn page_locs(full_page: usize) -> Vec<i64> {
1697        let base = (full_page * P) as i64;
1698        (0..P as i64).map(|offset| base + offset).collect()
1699    }
1700
1701    /// A base that is not a multiple of the page unit puts two full pages
1702    /// in one ring block. The free list is kept in units precisely so
1703    /// that cannot be expressed.
1704    #[test]
1705    fn every_unit_base_is_a_multiple_of_the_page_unit() {
1706        let mut allocator = FreeListAllocator::new(8 * P, P);
1707        assert_eq!(allocator.available(), 8 * P);
1708        let taken = allocator.alloc(3).expect("three of eight");
1709        assert_eq!(taken.len(), 3);
1710        assert!(taken.iter().all(|base| base.is_multiple_of(P)), "{taken:?}");
1711        // LIFO from the tail, ascending inside the slice.
1712        assert_eq!(taken, vec![5 * P, 6 * P, 7 * P]);
1713        assert_eq!(allocator.available(), 5 * P);
1714
1715        allocator.free(&taken);
1716        assert_eq!(allocator.available(), 8 * P);
1717        // The just-freed pages come straight back.
1718        assert_eq!(allocator.alloc(1).unwrap(), vec![7 * P]);
1719    }
1720
1721    #[test]
1722    #[should_panic(expected = "is not a unit base")]
1723    fn returning_a_base_that_is_not_a_unit_base_is_refused() {
1724        let mut allocator = FreeListAllocator::new(8 * P, P);
1725        allocator.free(&[P + 1]);
1726    }
1727
1728    #[test]
1729    fn an_oversized_allocation_takes_nothing() {
1730        let mut allocator = FreeListAllocator::new(4 * P, P);
1731        let err = allocator.alloc(5).unwrap_err();
1732        assert_eq!(err.needed_units, 5);
1733        assert_eq!(err.available_units, 4);
1734        assert_eq!(allocator.available(), 4 * P, "nothing was taken");
1735        allocator.alloc(4).expect("the free list is intact");
1736    }
1737
1738    /// A refused window allocation must leave the mapping untouched: the
1739    /// caller's next move is to evict and retry, not to undo a partial
1740    /// binding.
1741    #[test]
1742    fn an_exhausted_window_pool_binds_nothing() {
1743        let mut pool = pool();
1744        let mut locs = Vec::new();
1745        for full_page in 0..5 {
1746            locs.extend(page_locs(full_page));
1747        }
1748        let err = pool.alloc_swa(&locs).unwrap_err();
1749        assert_eq!(err.needed_units, 5);
1750        assert_eq!(err.available_units, 4);
1751        assert_eq!(pool.translate(0), NO_WINDOW_SLOT);
1752        pool.check_integrity();
1753    }
1754
1755    /// The state ring's block layout is keyed on the in-page offset, so a
1756    /// binding that permutes inside its page lands on the wrong ring slot.
1757    #[test]
1758    fn alloc_swa_preserves_in_page_offsets() {
1759        let mut pool = pool();
1760        pool.alloc_swa(&page_locs(0)).expect("one of four");
1761        let base = pool.translate(0);
1762        assert!(base >= 0 && (base as usize).is_multiple_of(P));
1763        for offset in 0..P as i64 {
1764            assert_eq!(pool.translate(offset), base + offset);
1765        }
1766        pool.check_integrity();
1767    }
1768
1769    #[test]
1770    #[should_panic(expected = "whole pages")]
1771    fn alloc_swa_refuses_a_partial_page() {
1772        let mut pool = pool();
1773        let _ = pool.alloc_swa(&page_locs(0)[..P - 1]);
1774    }
1775
1776    #[test]
1777    #[should_panic(expected = "contiguous ascending")]
1778    fn alloc_swa_refuses_a_page_that_is_not_contiguous_ascending() {
1779        let mut pool = pool();
1780        let mut locs = page_locs(0);
1781        locs.swap(3, 9);
1782        let _ = pool.alloc_swa(&locs);
1783    }
1784
1785    /// Freeing half a page returns a window page whose other half is
1786    /// still mapped: the next allocation gets a page two full pages read
1787    /// through.
1788    #[test]
1789    #[should_panic(expected = "partial pages")]
1790    fn free_swa_refuses_partial_pages() {
1791        let mut pool = pool();
1792        pool.alloc_swa(&page_locs(0)).unwrap();
1793        pool.free_swa(&page_locs(0)[..P - 1]);
1794    }
1795
1796    #[test]
1797    fn freeing_the_same_page_twice_is_a_no_op() {
1798        let mut pool = pool();
1799        pool.alloc_swa(&page_locs(1)).unwrap();
1800        assert_eq!(pool.swa_available_size(), 3 * P);
1801        pool.free_swa(&page_locs(1));
1802        assert_eq!(pool.swa_available_size(), 4 * P);
1803        assert_eq!(pool.translate(P as i64), NO_WINDOW_SLOT);
1804        // A slide, a tombstone and an eviction pass may all name it.
1805        pool.free_swa(&page_locs(1));
1806        assert_eq!(
1807            pool.swa_available_size(),
1808            4 * P,
1809            "the second free took nothing"
1810        );
1811        pool.free_swa(&page_locs(2));
1812        assert_eq!(pool.swa_available_size(), 4 * P);
1813        pool.check_integrity();
1814    }
1815
1816    #[test]
1817    fn a_negative_or_unbound_loc_translates_to_the_sentinel() {
1818        let pool = pool();
1819        assert_eq!(pool.translate(-1), NO_WINDOW_SLOT);
1820        assert_eq!(pool.translate(-99), NO_WINDOW_SLOT);
1821        assert_eq!(pool.translate(0), NO_WINDOW_SLOT);
1822        // The trailing sentinel row is addressable and permanently -1.
1823        assert_eq!(pool.translate((8 * P) as i64), NO_WINDOW_SLOT);
1824    }
1825
1826    /// The dummy page is bound so graph-padded rows scatter to a real
1827    /// slot, and excluded from the free list so it can never be handed to
1828    /// a request.
1829    #[test]
1830    fn the_dummy_page_is_bound_outside_the_free_list() {
1831        let pool = pool();
1832        assert_eq!(pool.translate((7 * P) as i64), (4 * P) as i64);
1833        assert_eq!(pool.swa_available_size(), 4 * P);
1834        assert_eq!(pool.swa_num_tokens(), 4 * P + 1);
1835        pool.check_integrity();
1836    }
1837
1838    /// The cap is halved because a batched prefill holds a whole chunk's
1839    /// window live at once, and the concurrent working set comes off
1840    /// first.
1841    #[test]
1842    fn the_prefill_chunk_cap_halves_what_is_left_after_the_reserve() {
1843        let sizes = dsv4_pool_sizes(128, &args(), 0.5, P, Some(41));
1844        let pool = Dsv4WindowPool::new(&sizes, 2, true);
1845        let reserved = dsv4_reserved_window_pages(2, true); // 13
1846        assert_eq!(pool.prefill_chunk_budget(), (40 - reserved) / 2 * P);
1847        // A pool that cannot cover its own reserve still admits one page.
1848        assert_eq!(pool.prefill_chunk_budget(), 13 * P);
1849        let tiny = Dsv4WindowPool::new(&dsv4_pool_sizes(8, &args(), 0.5, P, Some(5)), 2, true);
1850        assert_eq!(tiny.prefill_chunk_budget(), P);
1851    }
1852
1853    /// `ring_size | P`, so two window pages never share a ring block --
1854    /// and the location is DERIVED from the live binding. Storing it is
1855    /// the naive way, and this shows what it costs: once a window page is
1856    /// recycled, a `state_loc` cached from its previous owner addresses
1857    /// the new owner's carry state, silently.
1858    #[test]
1859    fn state_loc_is_derived_so_two_pages_never_share_a_ring_block() {
1860        let ring = 8;
1861        // Page 0 covers ring block 0..8, page 1 covers 8..16.
1862        for offset in 0..P as i64 {
1863            assert_eq!(Dsv4WindowPool::state_loc(offset, ring, P), offset % 8);
1864            assert_eq!(
1865                Dsv4WindowPool::state_loc(P as i64 + offset, ring, P),
1866                8 + offset % 8
1867            );
1868        }
1869        assert_eq!(Dsv4WindowPool::state_loc(NO_WINDOW_SLOT, ring, P), -1);
1870
1871        let mut pool = pool();
1872        pool.alloc_swa(&page_locs(0)).unwrap();
1873        pool.alloc_swa(&page_locs(1)).unwrap();
1874        // What a caller that STORED the location would keep holding.
1875        let stored = Dsv4WindowPool::state_loc(pool.translate(0), ring, P);
1876
1877        pool.free_swa(&page_locs(0));
1878        pool.alloc_swa(&page_locs(2)).unwrap();
1879        let recycled = Dsv4WindowPool::state_loc(pool.translate((2 * P) as i64), ring, P);
1880        assert_eq!(
1881            stored, recycled,
1882            "the recycled page inherits the ring block, so a stored state_loc reads its carry"
1883        );
1884        // Derived from the live binding, page 0 now has no state at all.
1885        assert_eq!(
1886            Dsv4WindowPool::state_loc(pool.translate(0), ring, P),
1887            NO_WINDOW_SLOT
1888        );
1889        pool.check_integrity();
1890    }
1891
1892    #[test]
1893    #[should_panic(expected = "must divide")]
1894    fn a_ring_that_does_not_divide_the_page_is_refused() {
1895        Dsv4WindowPool::state_loc(0, 7, P);
1896    }
1897
1898    /// `p = pos - ((pos - j) % win)` names the position in ring slot `j`.
1899    /// The modulo is euclidean: a truncated remainder would name a
1900    /// position AFTER `pos` for `j > pos`, read as though the request had
1901    /// already written it.
1902    #[test]
1903    fn the_ring_names_the_latest_position_congruent_to_each_slot() {
1904        let win = 4;
1905        // Mid-sequence: every slot holds one of the last `win` positions.
1906        let held: Vec<i64> = (0..win)
1907            .map(|j| window_ring_position(10, j, win).unwrap())
1908            .collect();
1909        assert_eq!(held, vec![8, 9, 10, 7]);
1910
1911        // Early decode: the slots the sequence has not reached are masked.
1912        assert_eq!(window_ring_position(1, 0, win), Some(0));
1913        assert_eq!(window_ring_position(1, 1, win), Some(1));
1914        assert_eq!(window_ring_position(1, 2, win), None);
1915        assert_eq!(window_ring_position(1, 3, win), None);
1916        // A truncated `%` would have said 1 - ((1 - 3) % 4) == 3 here.
1917        assert_ne!(window_ring_position(1, 3, win), Some(3));
1918    }
1919
1920    #[test]
1921    fn the_window_context_reads_the_ring_through_the_live_mapping() {
1922        let mut pool = pool();
1923        pool.alloc_swa(&page_locs(0)).unwrap();
1924        let full_locs: Vec<i64> = (0..P as i64).collect();
1925
1926        let ctx = pool.window_ctx(3, &full_locs);
1927        assert_eq!(ctx.window_slot, pool.translate(3));
1928        assert_eq!(ctx.prev_window_slot, pool.translate(2));
1929        assert_eq!(ctx.window_slots_topk.len(), P);
1930        // Only the first four slots have been reached.
1931        for (j, slot) in ctx.window_slots_topk.iter().enumerate() {
1932            let expected = if j <= 3 {
1933                pool.translate(j as i64)
1934            } else {
1935                NO_WINDOW_SLOT
1936            };
1937            assert_eq!(*slot, expected, "ring slot {j}");
1938        }
1939
1940        // Position 0 has no predecessor: the clamp reads itself, never -1.
1941        let first = pool.window_ctx(0, &full_locs);
1942        assert_eq!(first.prev_window_slot, pool.translate(0));
1943    }
1944
1945    /// A long-running page-granular workload must conserve exactly,
1946    /// however many times the window slides.
1947    #[test]
1948    fn the_window_pool_conserves_every_page() {
1949        let sizes = dsv4_pool_sizes(64, &args(), 0.5, P, Some(9));
1950        let mut pool = Dsv4WindowPool::new(&sizes, 1, true);
1951        let live_pages = 4;
1952        for full_page in 0..48 {
1953            pool.alloc_swa(&page_locs(full_page % 60))
1954                .unwrap_or_else(|err| panic!("page {full_page}: {err}"));
1955            if full_page >= live_pages {
1956                pool.free_swa(&page_locs((full_page - live_pages) % 60));
1957            }
1958            pool.check_integrity();
1959        }
1960        // Exactly the live window is held: 8 allocatable pages, 4 bound.
1961        assert_eq!(pool.swa_available_size(), (8 - live_pages) * P);
1962    }
1963
1964    /// The equality is what catches a leak; a `<=` would tolerate it.
1965    #[test]
1966    #[should_panic(expected = "leaked or double-freed")]
1967    fn the_invariant_catches_a_leaked_window_page() {
1968        let mut pool = pool();
1969        pool.alloc_swa(&page_locs(0)).unwrap();
1970        // A caller that dropped the binding without handing the page back.
1971        pool.unbind_window_pages(&page_locs(0));
1972        pool.check_integrity();
1973    }
1974
1975    /// And a window page bound to two full pages -- the aliasing the
1976    /// page-atomic free list exists to prevent.
1977    #[test]
1978    #[should_panic(expected = "bound to two full pages")]
1979    fn the_invariant_catches_a_window_page_bound_twice() {
1980        let mut pool = pool();
1981        pool.alloc_swa(&page_locs(0)).unwrap();
1982        let window_base = pool.translate(0) as usize;
1983        pool.bind_window_pages(P, window_base);
1984        pool.check_integrity();
1985    }
1986}