metrale_model_engine/traits/model/lifecycle.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `ModelLifecycle`, one of the supertraits `Model` is made of. Its methods, default
4//! bodies and docs are the ones `Model` declared before the split.
5//!
6//! Owner: model-engine.
7//! Invariants: none beyond the types.
8
9use crate::traits::SequenceState;
10use anyhow::{Result, bail};
11
12/// 2026-09-30: What [`ModelLifecycle::reclaim_prefix_blocks_for`] left: the free-block
13/// count a prefill of the prompt needs (`target`) and the count reached (`free`).
14/// `free < target` means nothing more was evictable.
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub struct PrefixReclaim {
17 pub target: usize,
18 pub free: usize,
19}
20
21/// 2026-09-26: The model's lifecycle and memory: teardown, thread binding, sequence allocation and
22/// release, sequence swap, and the KV pool counts.
23pub trait ModelLifecycle {
24 /// 2026-09-25: Release the device memory this model owns, newest first. The scheduler calls
25 /// it once at shutdown (`core/finish.rs`). Default: `Ok(())`, for models that own no pooled
26 /// device memory; a model that owns pools and keeps the default leaks them.
27 fn teardown(&mut self) -> Result<()> {
28 Ok(())
29 }
30
31 /// 2026-09-25: Poll this model's TQ+ InnerQ calibration driver; the scheduler calls it once
32 /// per prefill step. Default: no-op.
33 fn poll_innerq(&self) {}
34
35 /// 2026-09-25: Model dimensions for the `--high-speed-swap` orchestrator; `None` for a model
36 /// that does not support it (the default).
37 fn high_speed_swap_dims(&self) -> Option<metrale_storage::ModelDims> {
38 None
39 }
40
41 /// 2026-09-25: Bind the GPU context to the current thread. The scheduler thread and the EP
42 /// worker call it before using the model.
43 fn bind_gpu_to_thread(&self) -> Result<()>;
44
45 /// 2026-09-25: A new sequence with its layer states.
46 fn alloc_sequence(&self) -> Result<SequenceState>;
47
48 /// 2026-09-25: [`Self::alloc_sequence`] for a request that can reach at most `budget_tokens`
49 /// (`prompt_len + max_tokens`), so proposer state that scales with context can be sized to
50 /// that (`DraftProposer::alloc_state_for`). Default: [`Self::alloc_sequence`].
51 fn alloc_sequence_for(&self, budget_tokens: usize) -> Result<SequenceState> {
52 let _ = budget_tokens;
53 self.alloc_sequence()
54 }
55
56 /// 2026-09-25: Insert the sequence's tokens (prompt and generated) into the prefix cache. Call
57 /// before [`Self::free_sequence`], while its blocks are still valid.
58 fn cache_sequence(&self, seq: &SequenceState);
59
60 /// 2026-09-25: Release the sequence's KV blocks, prefix-cache refs, SSM pool slot and LoRA
61 /// slot ref.
62 fn free_sequence(&self, seq: &mut SequenceState) -> Result<()>;
63
64 /// 2026-09-25: Move the sequence's SSM state (`h_state` and `conv_state` of every SSM layer)
65 /// to pool slot `new_slot`.
66 fn compact_sequence(&self, seq: &mut SequenceState, new_slot: usize) -> Result<()>;
67
68 /// 2026-09-25: Give up a retired sequence's SSM pool slot after `compact_sequence` moved a
69 /// surviving sequence into it: sets `slot_idx` to `usize::MAX` and empties the slot guard, so
70 /// freeing or dropping this sequence does not release the slot. Call it right after that
71 /// `compact_sequence`, before any step that can drop the sequence.
72 fn detach_slot_for_reuse(&self, seq: &mut SequenceState);
73
74 /// 2026-09-25: Write the sequence's KV blocks and SSM state to `writer`, in a format the
75 /// model owns, without freeing anything. Default: an error.
76 fn save_sequence_state(
77 &self,
78 _seq: &SequenceState,
79 _writer: &mut dyn std::io::Write,
80 ) -> Result<()> {
81 bail!("swap not supported by this model")
82 }
83
84 /// 2026-09-25: Read state written by [`Self::save_sequence_state`] into an allocated
85 /// sequence, allocating `num_blocks` KV blocks. Default: an error.
86 fn restore_sequence_state(
87 &self,
88 _seq: &mut SequenceState,
89 _num_blocks: usize,
90 _reader: &mut dyn std::io::Read,
91 ) -> Result<()> {
92 bail!("swap not supported by this model")
93 }
94
95 /// 2026-09-25: Free KV blocks. Default `0`.
96 fn num_free_blocks(&self) -> usize {
97 0
98 }
99
100 /// 2026-09-25: Total KV blocks in the paged cache. Default `0`.
101 fn num_total_blocks(&self) -> usize {
102 0
103 }
104
105 /// 2026-10-01: Bytes one block index of the paged cache holds across every layer, K and V,
106 /// so `num_total_blocks() * kv_block_bytes()` is the pool's size. `0` for a model without a
107 /// paged cache. No default: a model that has a pool must say what its blocks cost.
108 fn kv_block_bytes(&self) -> usize;
109
110 /// 2026-09-25: Evict up to `num_blocks` blocks from the prefix cache; returns how many became
111 /// free. Swap-in and preemption resume check `num_free_blocks()` before they allocate, so
112 /// they call this to reach blocks the cache holds. `0` means nothing was evictable, and the
113 /// caller stops asking. Default `0`.
114 fn reclaim_prefix_blocks(&self, _num_blocks: usize) -> usize {
115 0
116 }
117
118 /// 2026-09-30: Reclaim prefix-cache blocks until `num_free_blocks()` covers what a prefill
119 /// of `prompt` under `adapter_slot` allocates: `blocks_needed` (the caller's count for the
120 /// whole prompt, headroom included) less the blocks of the cached prefix that prefill will
121 /// reuse. An empty `prompt` means the prefill reuses nothing (prompt logprobs). Admission
122 /// and preemption resume call it before they allocate. Default: no prefix is reused, and
123 /// [`Self::reclaim_prefix_blocks`] runs until `blocks_needed` is free or it returns 0.
124 fn reclaim_prefix_blocks_for(
125 &self,
126 prompt: &[u32],
127 adapter_slot: i32,
128 blocks_needed: usize,
129 ) -> PrefixReclaim {
130 let _ = (prompt, adapter_slot);
131 let mut free = self.num_free_blocks();
132 while free < blocks_needed && self.reclaim_prefix_blocks(blocks_needed - free) > 0 {
133 free = self.num_free_blocks();
134 }
135 PrefixReclaim {
136 target: blocks_needed,
137 free,
138 }
139 }
140}