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}