Support stateful KV cache for multimodal decoder

src/models/multi_modal.cpp

@@ -662,6 +662,7 @@ DecoderState::DecoderState(const MultiModalLanguageModel& model, DeviceSpan<int3
     : State{params, model},
       model_{model},
       position_inputs_{CreatePositionInputs(*this, sequence_lengths, model_.config_->model.decoder.inputs.attention_mask)},
+      kv_cache_{CreateKeyValueCache(*this)},
       recurrent_state_{CreateRecurrentState(*this)} {
   inputs_embeds_.Add();
 
@@ -689,7 +690,8 @@ DecoderState::DecoderState(const MultiModalLanguageModel& model, DeviceSpan<int3
 
   position_inputs_->Add();
   logits_.Add();
-  kv_cache_.Add();
+  if (kv_cache_)
+    kv_cache_->Add();
   if (recurrent_state_)
     recurrent_state_->Add();
 }
@@ -749,7 +751,8 @@ DeviceSpan<float> DecoderState::RunPrefillWithChunking(int current_length, Devic
 
     if (decoder_input_ids_) decoder_input_ids_->Update(chunk_tokens);
     position_inputs_->Update(chunk_tokens, length, static_cast<int>(current_chunk_size));
-    kv_cache_.Update(next_indices, length);
+    if (kv_cache_)
+      kv_cache_->Update(next_indices, length);
     if (recurrent_state_)
       recurrent_state_->Update();
     logits_.Update(chunk_tokens, current_chunk_size);
@@ -776,7 +779,8 @@ void DecoderState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, int tot
   size_t new_length = next_tokens.size() / batch_size;
   if (decoder_input_ids_) decoder_input_ids_->Update(next_tokens);
   position_inputs_->Update(next_tokens, total_length, static_cast<int>(new_length));
-  kv_cache_.Update(beam_indices, total_length);
+  if (kv_cache_)
+    kv_cache_->Update(beam_indices, total_length);
   if (recurrent_state_)
     recurrent_state_->Update();
   logits_.Update(next_tokens, new_length);
@@ -787,14 +791,24 @@ void DecoderState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, int tot
 // Overload for pipeline to call
 void DecoderState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, int total_length, DeviceSpan<int32_t> beam_indices, size_t new_length) {
   if (decoder_input_ids_) decoder_input_ids_->Update(next_tokens);
-  kv_cache_.Update(beam_indices, total_length);
+  if (kv_cache_)
+    kv_cache_->Update(beam_indices, total_length);
   if (recurrent_state_)
     recurrent_state_->Update();
   logits_.Update(next_tokens, new_length);
   inputs_embeds_.UpdateSequenceLength(new_length);
   if (per_layer_inputs_) per_layer_inputs_->UpdateSequenceLength(new_length);
 }
 
+void DecoderState::RewindTo(size_t index) {
+  if (position_inputs_)
+    position_inputs_->RewindTo(index);
+  if (kv_cache_)
+    kv_cache_->RewindTo(index);
+  if (recurrent_state_)
+    recurrent_state_->RewindTo(index);
+}
+
 MultiModalPipelineState::MultiModalPipelineState(const MultiModalLanguageModel& model, DeviceSpan<int32_t> sequence_lengths, const GeneratorParams& params)
     : State{params, model},
       model_{model},
@@ -925,6 +939,13 @@ DeviceSpan<float> MultiModalPipelineState::Run(int current_length, DeviceSpan<in
   return decoder_state_->Run(current_length, next_tokens, next_indices);
 }
 
+void MultiModalPipelineState::RewindTo(size_t index) {
+  if (decoder_state_)
+    decoder_state_->RewindTo(index);
+  if (index == 0)
+    is_prompt_ = true;
+}
+
 OrtValue* MultiModalPipelineState::GetInput(const char* name) {
   if (vision_state_) {
     // Check if input name is in vision state's inputs

src/models/multi_modal.h

@@ -178,6 +178,7 @@ struct DecoderState : State {
 
   DeviceSpan<float> Run(int current_length, DeviceSpan<int32_t>& next_tokens, DeviceSpan<int32_t> next_indices) override;
   void UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, int current_length, DeviceSpan<int32_t> beam_indices);
+  void RewindTo(size_t index) override;
 
   // Prefill chunking (see search.chunk_size). The embedding model still runs once over the whole
   // prompt (it is a lookup/projection), while the decoder prefill is split into several runs so the
@@ -198,7 +199,7 @@ struct DecoderState : State {
   std::unique_ptr<Embeddings> per_layer_inputs_;        // Optional model input (Gemma4: per-layer conditioning)
   std::unique_ptr<DefaultInputIDs> decoder_input_ids_;  // Optional model input (e.g., Gemma4 decoder needs input_ids)
   std::unique_ptr<PositionInputs> position_inputs_;     // Model input
-  DefaultKeyValueCache kv_cache_{*this};                // Model input
+  std::unique_ptr<KeyValueCache> kv_cache_;             // Model input (ModelManaged for stateful models)
   std::unique_ptr<RecurrentState> recurrent_state_;     // Model input (for hybrid models)
   Logits logits_{*this};                                // Model output
 };
@@ -218,6 +219,8 @@ struct MultiModalPipelineState : State {
 
   OrtValue* GetOutput(const char* name) override;
 
+  void RewindTo(size_t index) override;
+
  private:
   void UpdateInputsOutputs(const DeviceSpan<int32_t>& next_tokens, DeviceSpan<int32_t> next_indices,
                            int current_length);