Add onStageComplete

src/models/decoder_only_pipeline.cpp

@@ -344,6 +344,10 @@ void DecoderOnlyPipelineState::RunPipeline(int total_length, DeviceSpan<int32_t>
         }
       }
     }
+
+    // Notify derived classes that this pipeline stage has completed.
+    // This allows e.g. Qwen VL to inject vision embeddings after the embeddings stage.
+    OnStageComplete(pipeline_state->id_);
   }
 }
 

src/models/decoder_only_pipeline.h

@@ -74,8 +74,7 @@ struct DecoderOnlyPipelineState : State {
   // Virtual hook called after each pipeline stage completes, before next stage starts.
   // Allows derived classes to modify stage outputs (e.g., inject vision embeddings).
   // stage_id: ID of the stage that just completed
-  // next_tokens: current input tokens for pipeline
-  virtual void OnStageComplete(size_t stage_id, DeviceSpan<int32_t>& next_tokens) {}
+  virtual void OnStageComplete(size_t stage_id) {}
 
   // Stores all the outputs from the previous pipeline state(s)
   std::unordered_map<std::string, std::unique_ptr<OrtValue>> ortvalue_store_;

src/models/qwen_vl_model.cpp

@@ -150,17 +150,16 @@ void Qwen2_5_VL_PipelineState::SetExtraInputs(const std::vector<ExtraInput>& ext
   vision_ran_ = true;
 }
 
-void Qwen2_5_VL_PipelineState::OnStageComplete(size_t stage_id, DeviceSpan<int32_t>& next_tokens) {
+void Qwen2_5_VL_PipelineState::OnStageComplete(size_t stage_id) {
   if (stage_id != 0 || !vision_ran_) return;
 
   const auto& embeddings_config = vl_model_.config_->model.decoder.pipeline[0];
   if (!embeddings_config.outputs.empty()) {
-    InjectVisionEmbeddings(embeddings_config.outputs[0], next_tokens);
+    InjectVisionEmbeddings(embeddings_config.outputs[0]);
   }
 }
 
-void Qwen2_5_VL_PipelineState::InjectVisionEmbeddings(const std::string& embeddings_output_name,
-                                                      DeviceSpan<int32_t>& input_token_ids) {
+void Qwen2_5_VL_PipelineState::InjectVisionEmbeddings(const std::string& embeddings_output_name) {
   auto it = ortvalue_store_.find(embeddings_output_name);
   if (it == ortvalue_store_.end() || !it->second) {
     throw std::runtime_error("Vision embedding injection: embeddings output '" + embeddings_output_name + "' not found in ortvalue_store");

src/models/qwen_vl_model.h

@@ -27,11 +27,10 @@ struct Qwen2_5_VL_PipelineState : public DecoderOnlyPipelineState {
   void SetExtraInputs(const std::vector<ExtraInput>& extra_inputs) override;
 
  protected:
-  void OnStageComplete(size_t stage_id, DeviceSpan<int32_t>& next_tokens) override;
+  void OnStageComplete(size_t stage_id) override;
 
  private:
-  void InjectVisionEmbeddings(const std::string& embeddings_output_name,
-                              DeviceSpan<int32_t>& input_token_ids);
+  void InjectVisionEmbeddings(const std::string& embeddings_output_name);
 
   const Qwen2_5_VL_PipelineModel& vl_model_;
   bool vision_ran_{false};