Make Sequence::set_toks more safe (#1190)

Author: EricLBuehler

mistralrs-core/src/sequence.rs

@@ -1,6 +1,6 @@
 use crate::{
     get_mut_group,
-    pipeline::LayerCaches,
+    pipeline::{text_models_inputs_processor::PagedAttentionMeta, LayerCaches},
     response::{ChatCompletionChunkResponse, Choice, ChunkChoice, Response, SYSTEM_FINGERPRINT},
     sampler::{Logprobs, Sampler},
     ChatCompletionResponse, Usage,
@@ -480,7 +480,11 @@ impl Sequence {
     }
 
     /// This will also set prompt_len
-    pub(crate) fn set_toks(&mut self, toks: Vec<u32>) {
+    pub(crate) fn set_toks_and_reallocate(
+        &mut self,
+        toks: Vec<u32>,
+        paged_attn_metadata: Option<&mut PagedAttentionMeta<'_>>,
+    ) {
         self.tokens.clone_from(&toks);
         self.prompt_len = self.tokens.len();
         // Handle possible block engine
@@ -495,6 +499,12 @@ impl Sequence {
         }
         self.custom_metadata
             .append_tokens_to_blocks(toks.iter().map(|x| *x as usize).collect::<Vec<_>>());
+
+        if let Some(metadata) = paged_attn_metadata {
+            // Free and then reallocate as appropriate
+            metadata.block_engine.free_sequence(*self.id());
+            metadata.block_engine.allocate(self);
+        }
     }
 
     pub fn completion_bytes(&self) -> &[u8] {

mistralrs-core/src/vision_models/idefics3/inputs_processor.rs

@@ -250,13 +250,8 @@ impl InputsProcessor for Idefics3ImageProcessor {
 
                 let ids = toks.get_ids().to_vec();
                 all_ids.push(ids.clone());
-                seq.set_toks(ids);
 
-                if let Some(ref mut metadata) = paged_attn_metadata {
-                    // Free and then reallocate as appropriate
-                    metadata.block_engine.free_sequence(*seq.id());
-                    metadata.block_engine.allocate(*seq);
-                }
+                seq.set_toks_and_reallocate(ids, paged_attn_metadata.as_mut());
             }
 
             let mut all_ids_new = Vec::new();

mistralrs-core/src/vision_models/llava/llava_inputs_processor.rs

@@ -262,17 +262,13 @@ impl InputsProcessor for LLaVAInputProcessor {
                 input_ids.extend(item);
             }
             // NOTE(EricLBuehler): Casting to u32 is fine, we don't care about the other toks
-            seq.set_toks(
+            seq.set_toks_and_reallocate(
                 input_ids
                     .iter()
                     .map(|x| if *x < 0 { 0u32 } else { *x as u32 })
                     .collect::<Vec<_>>(),
+                paged_attn_metadata.as_mut(),
             );
-            if let Some(ref mut metadata) = paged_attn_metadata {
-                // Free and then reallocate as appropriate
-                metadata.block_engine.free_sequence(*seq.id());
-                metadata.block_engine.allocate(*seq);
-            }
 
             toks.push(input_ids);
         }

mistralrs-core/src/vision_models/llava/llava_next_inputs_processor.rs

@@ -306,17 +306,13 @@ impl InputsProcessor for LLaVANextInputProcessor {
                 input_ids.extend(item);
             }
             // NOTE(EricLBuehler): Casting to u32 is fine, we don't care about the other toks
-            seq.set_toks(
+            seq.set_toks_and_reallocate(
                 input_ids
                     .iter()
                     .map(|x| if *x < 0 { 0u32 } else { *x as u32 })
                     .collect::<Vec<_>>(),
+                paged_attn_metadata.as_mut(),
             );
-            if let Some(ref mut metadata) = paged_attn_metadata {
-                // Free and then reallocate as appropriate
-                metadata.block_engine.free_sequence(*seq.id());
-                metadata.block_engine.allocate(*seq);
-            }
 
             toks.push(input_ids);
         }

mistralrs-core/src/vision_models/minicpmo/inputs_processor.rs

@@ -315,12 +315,7 @@ impl InputsProcessor for MiniCpmOImageProcessor {
                         .get_ids()
                         .to_vec();
 
-                    seq.set_toks(input_ids.clone());
-                    if let Some(ref mut metadata) = paged_attn_metadata {
-                        // Free and then reallocate as appropriate
-                        metadata.block_engine.free_sequence(*seq.id());
-                        metadata.block_engine.allocate(*seq);
-                    }
+                    seq.set_toks_and_reallocate(input_ids.clone(), paged_attn_metadata.as_mut());
 
                     let image_start_idx = input_ids
                         .iter()

mistralrs-core/src/vision_models/phi3/phi3_inputs_processor.rs

@@ -299,17 +299,13 @@ impl InputsProcessor for Phi3InputsProcessor {
             }
 
             // NOTE(EricLBuehler): Casting to u32 is fine, we don't care about the other toks
-            seq.set_toks(
+            seq.set_toks_and_reallocate(
                 input_ids
                     .iter()
                     .map(|x| if *x < 0 { 0u32 } else { *x as u32 })
                     .collect::<Vec<_>>(),
+                paged_attn_metadata.as_mut(),
             );
-            if let Some(ref mut metadata) = paged_attn_metadata {
-                // Free and then reallocate as appropriate
-                metadata.block_engine.free_sequence(*seq.id());
-                metadata.block_engine.allocate(*seq);
-            }
 
             toks.push(input_ids);
         }

mistralrs-core/src/vision_models/phi4/inputs_processor.rs

@@ -241,12 +241,13 @@ impl InputsProcessor for Phi4MMInputsProcessor {
                 .replace_all(&detokenized, IMAGE_SPECIAL_TOKEN)
                 .to_string();
 
-            seq.set_toks(
+            seq.set_toks_and_reallocate(
                 tokenizer
                     .encode(detokenized.clone(), true)
                     .expect("Encode failed")
                     .get_ids()
                     .to_vec(),
+                paged_attn_metadata.as_mut(),
             );
 
             seq.set_initial_prompt(detokenized);
@@ -265,16 +266,10 @@ impl InputsProcessor for Phi4MMInputsProcessor {
                 let mut new_ids = seq.get_toks()[..i].to_vec();
                 new_ids.extend(vec![token_id; *token_count]);
                 new_ids.extend(seq.get_toks()[i + 1..].to_vec());
-                seq.set_toks(new_ids);
+                seq.set_toks_and_reallocate(new_ids, paged_attn_metadata.as_mut());
                 i += token_count;
             }
             toks.push(seq.get_toks().to_vec());
-
-            if let Some(ref mut metadata) = paged_attn_metadata {
-                // Free and then reallocate as appropriate
-                metadata.block_engine.free_sequence(*seq.id());
-                metadata.block_engine.allocate(*seq);
-            }
         }
 
         let iter = if is_prompt {

mistralrs-core/src/vision_models/qwen2vl/inputs_processor.rs

@@ -393,13 +393,7 @@ impl InputsProcessor for Qwen2VLImageProcessor {
                 let continuous_vid_pad = find_sequences(&ids, vid_pad[0]);
                 all_continuous_vid_pad.push(continuous_vid_pad);
 
-                seq.set_toks(ids);
-
-                if let Some(ref mut metadata) = paged_attn_metadata {
-                    // Free and then reallocate as appropriate
-                    metadata.block_engine.free_sequence(*seq.id());
-                    metadata.block_engine.allocate(*seq);
-                }
+                seq.set_toks_and_reallocate(ids, paged_attn_metadata.as_mut());
             }
 
             let mut input_ids_searching = Vec::new();