metrale_model_engine/
prefill_plan.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-27: Where prefill chunks end, and which chunk ends carry an SSM snapshot.
4//!
5//! A later request's prefix match is block-floored, and a snapshot deeper than the
6//! match cannot be restored. So a non-last chunk ends on a KV block boundary, and the
7//! chunk that would span the tail split point ends there. A last chunk that spans the
8//! split point is split by `TransformerModel::prefill_chunk_dispatch` instead, inside
9//! one scheduler tick: 2026-10-01, by capturing the SSM state at the split point inside
10//! its pass where the model and build support it, else as two passes.
11//!
12//! Before 2026-09-27 only the last chunk was split. A dense serve's first chunk is
13//! `prefill_budget + max_batch_size` = 8193 tokens, so every later chunk end sat one
14//! token past a block boundary, and a 32772-token prompt's split point (32752) fell
15//! inside the second-to-last chunk: no tail snapshot was saved, and the byte-identical
16//! warm request restored at 24577 and replayed 8195 tokens (12.5 s instead of ~0.2 s).
17//!
18//! Owner: model-engine prefill (SSM prefix cache).
19//! Invariants:
20//! - `plan_chunk_len` returns a length in `1..=proposed` for a non-empty `proposed`,
21//!   and returns `proposed` unchanged for a last chunk.
22
23/// 2026-09-27: The tail split point of a `total`-token prompt: one block below the last
24/// block boundary strictly under `total`, or `None` when that is not above token 0. A
25/// warm next turn's block-floored match lands at or one block below that boundary, so a
26/// snapshot here is restorable by both.
27pub fn tail_split_point(total: usize, block_size: usize) -> Option<usize> {
28    let tb = metrale_gpu_runtime::ssm_tail_boundary(total, block_size)?;
29    let cut = tb - block_size;
30    (cut > 0).then_some(cut)
31}
32
33/// 2026-09-27: Length of the prefill chunk at `offset` of a `total`-token prompt, given
34/// the capped `proposed` length.
35///
36/// A last chunk (`offset + proposed >= total`) is returned unchanged. A non-last chunk
37/// ends at the last `block_size` boundary inside it (when that is past `offset`), then
38/// at `split` when it would still span it, and is finally rounded down to a multiple of
39/// 4 (the GDN WY4 rule for intermediate chunks). With chunk ends on block boundaries
40/// every later offset is block-aligned, so the rounding changes nothing when
41/// `block_size` is a multiple of 4.
42pub fn plan_chunk_len(
43    offset: usize,
44    total: usize,
45    proposed: usize,
46    block_size: Option<usize>,
47    split: Option<usize>,
48) -> usize {
49    if proposed == 0 || offset + proposed >= total {
50        return proposed;
51    }
52    let mut end = offset + proposed;
53    if let Some(bs) = block_size.filter(|&bs| bs > 0) {
54        let aligned = (end / bs) * bs;
55        if aligned > offset {
56            end = aligned;
57        }
58    }
59    if let Some(cut) = split
60        && offset < cut
61        && cut < end
62    {
63        end = cut;
64    }
65    let len = end - offset;
66    if len >= 4 { (len / 4) * 4 } else { len }
67}
68
69/// 2026-09-27: Whether a non-last chunk ending at `end_token` saves an SSM snapshot: a
70/// prompt-tail end (the last block boundary under `total`, or one block below it), or a
71/// block boundary whose block index is a multiple of `interval` (0 turns interval
72/// snapshots off). Block 0 never does.
73pub fn is_checkpoint_chunk_end(
74    end_token: usize,
75    total: usize,
76    block_size: usize,
77    interval: usize,
78) -> bool {
79    if block_size == 0 {
80        return false;
81    }
82    let end_block = end_token / block_size;
83    let on_interval = interval > 0 && end_block.is_multiple_of(interval);
84    end_block != 0 && (is_prompt_tail_end(end_token, total, block_size) || on_interval)
85}
86
87/// 2026-09-27: Whether a chunk end is a prompt-tail checkpoint: the last block boundary
88/// under `total`, or one block below it (`tail_split_point`). A warm next turn's
89/// block-floored match lands at one of the two.
90pub fn is_prompt_tail_end(end_token: usize, total: usize, block_size: usize) -> bool {
91    if block_size == 0 {
92        return false;
93    }
94    let tail = (total.saturating_sub(1) / block_size) * block_size;
95    end_token == tail || (tail >= block_size && end_token == tail - block_size)
96}
97
98#[cfg(test)]
99#[path = "prefill_plan_tests.rs"]
100mod tests;
metrale_model_engine/prefill_plan.rs

metrale_model_engine/
prefill_plan.rs
