From 7f717ed76883342480f13d693ba9e9b945ac8dd6 Mon Sep 17 00:00:00 2001
From: Scott McKay <Scott.McKay@microsoft.com>
Date: Tue, 19 May 2026 12:05:03 +1000
Subject: [PATCH] Fix nemotron leaks

---
 src/models/nemotron_speech.cpp | 34 +++++++++++++++++++++++++++-------
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/src/models/nemotron_speech.cpp b/src/models/nemotron_speech.cpp
index 2eaa01fdc2..879999f40a 100644
--- a/src/models/nemotron_speech.cpp
+++ b/src/models/nemotron_speech.cpp
@@ -423,10 +423,16 @@ void NemotronSpeechState::RunEncoder() {
   DeviceSpan<int32_t> dummy_tokens;
   encoder_state_->Run(0, dummy_tokens);
 
-  // Grab encoder outputs
+  // Grab encoder outputs. Each session.Run allocates a fresh OrtValue per output slot,
+  // so the State owns them and must release every slot it doesn't otherwise move out
+  // (otherwise the next Run overwrites the pointer and leaks the previous tensor).
   encoded_output_.reset(encoder_state_->outputs_[0]);
   encoder_state_->outputs_[0] = nullptr;
-  encoded_len_ = *encoder_state_->outputs_[1]->GetTensorData<int64_t>();
+  {
+    std::unique_ptr<OrtValue> encoded_len_owner{encoder_state_->outputs_[1]};
+    encoded_len_ = *encoded_len_owner->GetTensorData<int64_t>();
+    encoder_state_->outputs_[1] = nullptr;
+  }
 
   // Cache outputs are already moved into encoder_state_->cache_ by Run()
   // But State::Run uses output pointers differently - the outputs are written to by ORT
@@ -472,20 +478,26 @@ std::span<const int32_t> NemotronSpeechState::StepToken() {
     prediction_state_->Run(0, dummy_tokens);
 
     // Reshape decoder output for joiner: [1, dim] -> [1, 1, dim]
-    auto dec_out_shape = prediction_state_->outputs_[0]->GetTensorTypeAndShapeInfo()->GetShape();
+    // Take ownership of prediction output[0] so it's released this iteration
+    // (the next prediction Run will overwrite the slot and leak the old buffer otherwise).
+    std::unique_ptr<OrtValue> pred_out0_owner{prediction_state_->outputs_[0]};
+    prediction_state_->outputs_[0] = nullptr;
+    auto dec_out_shape = pred_out0_owner->GetTensorTypeAndShapeInfo()->GetShape();
     auto decoder_frame_shape = std::array<int64_t, 3>{1, 1, dec_out_shape[1]};
     auto dec_out_type = model_.session_info_.GetOutputDataType(nemotron_config_.dec_out_outputs);
     auto decoder_frame = OrtValue::CreateTensor(allocator, decoder_frame_shape, dec_out_type);
     ByteWrapTensor(*model_.p_device_, *decoder_frame)
-        .CopyFrom(ByteWrapTensor(*model_.p_device_, *prediction_state_->outputs_[0]));
+        .CopyFrom(ByteWrapTensor(*model_.p_device_, *pred_out0_owner));
 
     // Run joiner
     joiner_state_->SetInputFrames(encoder_frame_.get(), decoder_frame.get());
     joiner_state_->Run(0, dummy_tokens);
 
-    // Argmax over logits
-    const float* logits = joiner_state_->outputs_[0]->GetTensorData<float>();
-    auto logits_shape = joiner_state_->outputs_[0]->GetTensorTypeAndShapeInfo()->GetShape();
+    // Argmax over logits. Take ownership of joiner output[0] so it's released this iteration.
+    std::unique_ptr<OrtValue> joiner_out0_owner{joiner_state_->outputs_[0]};
+    joiner_state_->outputs_[0] = nullptr;
+    const float* logits = joiner_out0_owner->GetTensorData<float>();
+    auto logits_shape = joiner_out0_owner->GetTensorTypeAndShapeInfo()->GetShape();
     int total_logits = 1;
     for (auto d : logits_shape) total_logits *= static_cast<int>(d);
 
@@ -501,6 +513,14 @@ std::span<const int32_t> NemotronSpeechState::StepToken() {
     }
 
     if (best_token == nemotron_config_.blank_id) {
+      // Discard the prediction-network LSTM outputs from this step; the prior
+      // lstm_state_ is kept because no token was emitted. Without this the
+      // next prediction Run overwrites the slots and leaks these tensors.
+      std::unique_ptr<OrtValue>{prediction_state_->outputs_[1]}.reset();
+      prediction_state_->outputs_[1] = nullptr;
+      std::unique_ptr<OrtValue>{prediction_state_->outputs_[2]}.reset();
+      prediction_state_->outputs_[2] = nullptr;
+
       time_step_++;
       symbol_step_ = 0;
       continue;