Module prefill_plan

Module prefill_plan 

Source
Expand description

2026-09-27: Where prefill chunks end, and which chunk ends carry an SSM snapshot.

A later request’s prefix match is block-floored, and a snapshot deeper than the match cannot be restored. So a non-last chunk ends on a KV block boundary, and the chunk that would span the tail split point ends there. A last chunk that spans the split point is split by TransformerModel::prefill_chunk_dispatch instead, inside one scheduler tick: 2026-10-01, by capturing the SSM state at the split point inside its pass where the model and build support it, else as two passes.

Before 2026-09-27 only the last chunk was split. A dense serve’s first chunk is prefill_budget + max_batch_size = 8193 tokens, so every later chunk end sat one token past a block boundary, and a 32772-token prompt’s split point (32752) fell inside the second-to-last chunk: no tail snapshot was saved, and the byte-identical warm request restored at 24577 and replayed 8195 tokens (12.5 s instead of ~0.2 s).

Owner: model-engine prefill (SSM prefix cache). Invariants:

  • plan_chunk_len returns a length in 1..=proposed for a non-empty proposed, and returns proposed unchanged for a last chunk.

Functions§

is_checkpoint_chunk_end
2026-09-27: Whether a non-last chunk ending at end_token saves an SSM snapshot: a prompt-tail end (the last block boundary under total, or one block below it), or a block boundary whose block index is a multiple of interval (0 turns interval snapshots off). Block 0 never does.
is_prompt_tail_end
2026-09-27: Whether a chunk end is a prompt-tail checkpoint: the last block boundary under total, or one block below it (tail_split_point). A warm next turn’s block-floored match lands at one of the two.
plan_chunk_len
2026-09-27: Length of the prefill chunk at offset of a total-token prompt, given the capped proposed length.
tail_split_point
2026-09-27: The tail split point of a total-token prompt: one block below the last block boundary strictly under total, or None when that is not above token 0. A warm next turn’s block-floored match lands at or one block below that boundary, so a snapshot here is restorable by both.