metrale_model_engine/model/trait_impl/
impl_lifecycle.rs1use anyhow::Result;
10
11use crate::model::types::TransformerModel;
12use crate::traits::{ModelLifecycle, SequenceState};
13
14impl ModelLifecycle for TransformerModel {
15 fn teardown(&mut self) -> Result<()> {
16 self.gdn_carry_flush_pending()?;
17 self.release_pools()
18 }
19
20 #[cfg(feature = "cuda")]
23 fn poll_innerq(&self) {
24 if let Some(driver) = self.innerq.as_ref()
25 && let Err(e) = driver.maybe_finalize(128)
26 {
27 tracing::warn!(target: "metrale_model_engine::model::trait_impl", "InnerQ maybe_finalize failed: {e:#}");
28 }
29 }
30
31 fn high_speed_swap_dims(&self) -> Option<metrale_storage::ModelDims> {
32 self.high_speed_swap_dims_dispatch()
33 }
34
35 fn bind_gpu_to_thread(&self) -> Result<()> {
36 self.bind_gpu_to_thread_dispatch()
37 }
38
39 fn alloc_sequence(&self) -> Result<SequenceState> {
40 self.gdn_carry_flush_pending()?;
41 self.alloc_sequence_dispatch(usize::MAX)
42 }
43
44 fn alloc_sequence_for(&self, budget_tokens: usize) -> Result<SequenceState> {
45 self.gdn_carry_flush_pending()?;
46 self.alloc_sequence_dispatch(budget_tokens)
47 }
48
49 fn cache_sequence(&self, seq: &SequenceState) {
50 if let Err(e) = self.gdn_carry_flush_pending() {
51 tracing::error!("gdn carry flush before {}: {e:#}", "cache_sequence");
52 }
53 self.cache_sequence_dispatch(seq)
54 }
55
56 fn free_sequence(&self, seq: &mut SequenceState) -> Result<()> {
57 self.gdn_carry_flush_pending()?;
58 self.free_sequence_dispatch(seq)
59 }
60
61 fn compact_sequence(&self, seq: &mut SequenceState, new_slot: usize) -> Result<()> {
62 self.gdn_carry_flush_pending()?;
63 self.compact_sequence_dispatch(seq, new_slot)
64 }
65
66 fn detach_slot_for_reuse(&self, seq: &mut SequenceState) {
67 if let Err(e) = self.gdn_carry_flush_pending() {
68 tracing::error!("gdn carry flush before {}: {e:#}", "detach_slot_for_reuse");
69 }
70 self.detach_slot_for_reuse_dispatch(seq)
71 }
72
73 fn save_sequence_state(
74 &self,
75 seq: &SequenceState,
76 writer: &mut dyn std::io::Write,
77 ) -> Result<()> {
78 self.gdn_carry_flush_pending()?;
79 self.save_sequence_state_dispatch(seq, writer)
80 }
81
82 fn restore_sequence_state(
83 &self,
84 seq: &mut SequenceState,
85 num_blocks: usize,
86 reader: &mut dyn std::io::Read,
87 ) -> Result<()> {
88 self.gdn_carry_flush_pending()?;
89 self.restore_sequence_state_dispatch(seq, num_blocks, reader)
90 }
91
92 fn num_free_blocks(&self) -> usize {
93 self.num_free_blocks_dispatch()
94 }
95
96 fn num_total_blocks(&self) -> usize {
97 self.num_total_blocks_dispatch()
98 }
99
100 fn kv_block_bytes(&self) -> usize {
101 self.kv_block_bytes_dispatch()
102 }
103
104 fn reclaim_prefix_blocks(&self, num_blocks: usize) -> usize {
105 self.reclaim_prefix_blocks_dispatch(num_blocks)
106 }
107
108 fn reclaim_prefix_blocks_for(
109 &self,
110 prompt: &[u32],
111 adapter_slot: i32,
112 blocks_needed: usize,
113 ) -> crate::traits::PrefixReclaim {
114 self.reclaim_prefix_blocks_for_dispatch(prompt, adapter_slot, blocks_needed)
115 }
116}