Whisper chanegs

src/models/model.cpp

@@ -394,10 +394,20 @@ void ConfigureNvTensorRtRtxProfile(const Config& config, OrtSessionOptions& sess
   // Get max context length from config
   const int max_context_len = config.model.context_length;
 
-  // Extract KV cache name patterns from decoder config
+  // Extract KV cache name patterns from decoder config (self-attention)
   std::string_view past_key_pattern = config.model.decoder.inputs.past_key_names;
   std::string_view past_value_pattern = config.model.decoder.inputs.past_value_names;
 
+  // Extract cross-attention KV cache name patterns (for encoder-decoder models like Whisper)
+  std::string_view cross_past_key_pattern = config.model.decoder.inputs.cross_past_key_names;
+  std::string_view cross_past_value_pattern = config.model.decoder.inputs.cross_past_value_names;
+  const bool has_cross_attention = !cross_past_key_pattern.empty() && !cross_past_value_pattern.empty();
+
+  // Cross-attention uses fixed encoder sequence length (e.g., 1500 for Whisper = 3000 frames / 2)
+  // This is the encoder output sequence length, which is constant for a given model
+  constexpr int whisper_encoder_seq_len = 1500;  // Whisper: 30 seconds * 100 fps / 2 = 1500
+  const int cross_attention_seq_len = (config.model.type == "whisper") ? whisper_encoder_seq_len : 0;
+
   // Helper function to add KV cache with sequence length
   const auto add_key_value_cache_shapes = [](std::ostringstream& shapes,
                                              int batch_size,
@@ -473,6 +483,21 @@ void ConfigureNvTensorRtRtxProfile(const Config& config, OrtSessionOptions& sess
     add_generation_input_shapes(max_shapes, batch_size, max_context_len);
     add_key_value_cache_shapes(max_shapes, batch_size, past_key_pattern, past_value_pattern, max_context_len - 1, num_layers, num_kv_heads, head_dim);
 
+    // Add cross-attention cache shapes if present (encoder-decoder models like Whisper)
+    // Cross-attention caches have FIXED encoder sequence length (doesn't change during generation)
+    if (has_cross_attention && cross_attention_seq_len > 0) {
+      // Cross-attention uses num_attention_heads (not num_kv_heads) because encoder K/V aren't grouped
+      const int cross_num_heads = config.model.decoder.num_attention_heads;
+
+      // Add cross-attention cache shapes to all profiles (same fixed seq_len for all)
+      add_key_value_cache_shapes(min_shapes, batch_size, cross_past_key_pattern, cross_past_value_pattern,
+                                 cross_attention_seq_len, num_layers, cross_num_heads, head_dim);
+      add_key_value_cache_shapes(opt_shapes, batch_size, cross_past_key_pattern, cross_past_value_pattern,
+                                 cross_attention_seq_len, num_layers, cross_num_heads, head_dim);
+      add_key_value_cache_shapes(max_shapes, batch_size, cross_past_key_pattern, cross_past_value_pattern,
+                                 cross_attention_seq_len, num_layers, cross_num_heads, head_dim);
+    }
+
     // Add the constructed profiles to session options
     session_options.AddConfigEntry("ep.nvtensorrtrtxexecutionprovider.nv_profile_min_shapes", min_shapes.str().c_str());
     session_options.AddConfigEntry("ep.nvtensorrtrtxexecutionprovider.nv_profile_opt_shapes", opt_shapes.str().c_str());
@@ -501,6 +526,22 @@ void ConfigureNvTensorRtRtxProfile(const Config& config, OrtSessionOptions& sess
                << Config::Defaults::AttentionMaskName << ":" << batch_size << "x" << max_context_len;
     add_key_value_cache_shapes(max_shapes, batch_size, past_key_pattern, past_value_pattern, max_context_len, num_layers, num_kv_heads, head_dim);
 
+    // Add cross-attention cache shapes if present (encoder-decoder models like Whisper)
+    // Cross-attention caches have FIXED encoder sequence length (doesn't change during generation)
+    if (has_cross_attention && cross_attention_seq_len > 0) {
+      // Cross-attention uses num_attention_heads (not num_kv_heads) because encoder K/V aren't grouped
+      const int cross_num_heads = config.model.decoder.num_attention_heads;
+
+      // Add cross-attention cache shapes to all profiles (same fixed seq_len for all)
+      // Note: For single-profile, min uses min_batch_size, opt uses opt_batch_size, max uses batch_size
+      add_key_value_cache_shapes(min_shapes, min_batch_size, cross_past_key_pattern, cross_past_value_pattern,
+                                 cross_attention_seq_len, num_layers, cross_num_heads, head_dim);
+      add_key_value_cache_shapes(opt_shapes, opt_batch_size, cross_past_key_pattern, cross_past_value_pattern,
+                                 cross_attention_seq_len, num_layers, cross_num_heads, head_dim);
+      add_key_value_cache_shapes(max_shapes, batch_size, cross_past_key_pattern, cross_past_value_pattern,
+                                 cross_attention_seq_len, num_layers, cross_num_heads, head_dim);
+    }
+
     // Add the constructed profiles to session options
     session_options.AddConfigEntry("ep.nvtensorrtrtxexecutionprovider.nv_profile_min_shapes", min_shapes.str().c_str());
     session_options.AddConfigEntry("ep.nvtensorrtrtxexecutionprovider.nv_profile_opt_shapes", opt_shapes.str().c_str());

src/models/whisper.cpp

@@ -1,10 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+#include <cassert>
 #include "../generators.h"
 #include "whisper.h"
-
 namespace Generators {
-
 WhisperModel::WhisperModel(std::unique_ptr<Config> config, OrtEnv& ort_env)
     : Model{std::move(config)} {
   encoder_session_options_ = OrtSessionOptions::Create();
@@ -85,10 +84,127 @@ WhisperDecoderState::WhisperDecoderState(const WhisperModel& model, const Genera
     ByteWrapTensor(*model_.p_device_inputs_, *cache_indirection_).Zero();
   }
 
+  // attention_mask is only supported for NvTensorRtRtx EP
+  const bool is_trt_rtx = model_.p_device_inputs_->GetType() == DeviceType::NvTensorRtRtx;
+  has_attention_mask_input_ = is_trt_rtx && HasAttentionMaskInput();
+  has_position_ids_input_ = is_trt_rtx && model_.session_info_.HasInput(model_.config_->model.decoder.inputs.position_ids);
+  // Note: attention_mask/position_ids are initialized lazily on first run when current_length is known
+
   output_cross_qk_name_ = ComposeKeyValueName(model_.config_->model.decoder.outputs.output_cross_qk_names, 0);
 }
 
+void WhisperDecoderState::CreateAndInitializeAttentionMask(int64_t valid_length) {
+  int64_t batch_size = attention_mask_shape_[0];
+  int64_t buffer_length = attention_mask_shape_[1];  // max_length for static, current length for dynamic
+
+  if (model_.p_device_inputs_->GetType() != DeviceType::NvTensorRtRtx) {
+    assert(false && "Whisper attention_mask handling is only supported for TRT-RTX EP.");
+    throw std::runtime_error("Whisper attention_mask handling requires TRT-RTX EP.");
+  }
+
+  auto& gpu_allocator = model_.p_device_inputs_->GetAllocator();
+  attention_mask_ = OrtValue::CreateTensor(gpu_allocator, attention_mask_shape_, mask_type_);
+  auto attention_mask_span = ByteWrapTensor(*model_.p_device_inputs_, *attention_mask_);
+  attention_mask_span.Zero();
+
+  if (valid_length > 0) {
+    const bool updated_on_device = model_.p_device_inputs_->UpdateAttentionMask(nullptr,
+                                                                                 attention_mask_->GetTensorMutableRawData(),
+                                                                                 static_cast<int>(batch_size),
+                                                                                 static_cast<int>(valid_length),
+                                                                                 static_cast<int>(valid_length),
+                                                                                 static_cast<int>(buffer_length),
+                                                                                 true,
+                                                                                 mask_type_);
+    if (!updated_on_device) {
+      assert(false && "TRT-RTX attention_mask initialization must remain on GPU.");
+      throw std::runtime_error("TRT-RTX attention_mask initialization unexpectedly attempted CPU fallback.");
+    }
+  }
+}
+
+void WhisperDecoderState::UpdateAttentionMask(int current_length, int new_kv_length) {
+  int64_t batch_size = attention_mask_shape_[0];
+
+  if (use_static_buffer_) {
+    // Static buffer mode: update in-place
+    int64_t max_length = attention_mask_shape_[1];
+    const bool updated_on_device = model_.p_device_inputs_->UpdateAttentionMask(nullptr,
+                                                                                 attention_mask_->GetTensorMutableRawData(),
+                                                                                 static_cast<int>(batch_size),
+                                                                                 new_kv_length,
+                                                                                 current_length,
+                                                                                 static_cast<int>(max_length),
+                                                                                 true,
+                                                                                 mask_type_);
+    if (!updated_on_device) {
+      assert(false && "TRT-RTX static attention_mask updates must remain on GPU.");
+      throw std::runtime_error("TRT-RTX static attention_mask update unexpectedly attempted CPU fallback.");
+    }
+    return;
+  }
+
+  // Dynamic mode: create new tensor with expanded length
+  attention_mask_shape_[1] = current_length;
+  attention_mask_next_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), attention_mask_shape_, mask_type_);
+
+  const bool updated_on_device = model_.p_device_inputs_->UpdateAttentionMask(attention_mask_next_->GetTensorMutableRawData(),
+                                                                               attention_mask_->GetTensorMutableRawData(),
+                                                                               static_cast<int>(batch_size),
+                                                                               new_kv_length,
+                                                                               current_length,
+                                                                               params_->search.max_length,
+                                                                               false,
+                                                                               mask_type_);
+  if (!updated_on_device) {
+    assert(false && "TRT-RTX attention_mask updates must remain on GPU.");
+    throw std::runtime_error("TRT-RTX attention_mask update unexpectedly attempted CPU fallback.");
+  }
+
+  // Swap: next becomes current
+  std::swap(attention_mask_, attention_mask_next_);
+
+  // Update input pointer
+  for (size_t i = 0; i < input_names_.size(); ++i) {
+    if (input_names_[i] == model_.config_->model.decoder.inputs.attention_mask) {
+      inputs_[i] = attention_mask_.get();
+      break;
+    }
+  }
+}
+
 DeviceSpan<float> WhisperDecoderState::Run(int current_length, DeviceSpan<int32_t>& next_tokens, DeviceSpan<int32_t> next_indices) {
+  // Initialize attention_mask on first run (lazy initialization)
+  if (first_run_ && has_attention_mask_input_) {
+    // Get data type from model session
+    mask_type_ = model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.attention_mask);
+
+    // Validate type (must be INT32 or INT64)
+    if (mask_type_ != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 &&
+        mask_type_ != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) {
+      throw std::runtime_error("attention_mask must be int32 or int64 type");
+    }
+
+    // Check if using in-place KV cache (static buffer)
+    use_static_buffer_ = params_->IsPastPresentShareBufferEnabled(model_.config_->model.type);
+
+    // Set shape based on mode:
+    // - Static buffer (in-place KV cache): pre-allocate to max_length
+    // - Dynamic mode: allocate to current_length
+    if (use_static_buffer_) {
+      attention_mask_shape_ = {params_->BatchBeamSize(), params_->search.max_length};
+    } else {
+      attention_mask_shape_ = {params_->BatchBeamSize(), current_length};
+    }
+
+    // Create initial attention mask tensor
+    CreateAndInitializeAttentionMask(current_length);
+
+    // Register as input
+    input_names_.push_back(model_.config_->model.decoder.inputs.attention_mask.c_str());
+    inputs_.push_back(attention_mask_.get());
+  }
+
   // Add output QK on first run
   if (first_run_ && model_.session_info_.HasOutput(output_cross_qk_name_)) {
     output_cross_qk_type_ = model_.session_info_.GetOutputDataType(output_cross_qk_name_);
@@ -111,16 +227,70 @@ DeviceSpan<float> WhisperDecoderState::Run(int current_length, DeviceSpan<int32_
     State::SetRunOptions(model_.config_->model.decoder.run_options.value());
   }
   State::Run(*model_.session_decoder_);
+
   return logits_.Get();
 }
 
 void WhisperDecoderState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens, DeviceSpan<int32_t> beam_indices, int current_length, bool first_update) {
-  int batch_size = static_cast<int>(input_ids_.GetShape()[0]);
-  size_t new_length = next_tokens.size() / batch_size;
   input_ids_.Update(next_tokens);
+  const auto input_ids_shape = input_ids_.GetShape();
+  size_t new_length = static_cast<size_t>(input_ids_shape[1]);
   kv_cache_->Update(beam_indices, current_length);
   logits_.Update(next_tokens, first_run_ ? current_length : new_length);
 
+  // Update position_ids if model expects it
+  if (has_position_ids_input_) {
+    const auto& pos_name = model_.config_->model.decoder.inputs.position_ids;
+    if (position_ids_index_ == ~0U) {
+      position_ids_type_ = model_.session_info_.GetInputDataType(pos_name);
+      if (position_ids_type_ != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 &&
+          position_ids_type_ != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) {
+        throw std::runtime_error("position_ids must be int32 or int64 type");
+      }
+      position_ids_shape_ = {params_->BatchBeamSize(), static_cast<int64_t>(new_length)};
+      position_ids_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), position_ids_shape_, position_ids_type_);
+      position_ids_index_ = inputs_.size();
+      input_names_.push_back(pos_name.c_str());
+      inputs_.push_back(position_ids_.get());
+    } else if (position_ids_shape_[1] != static_cast<int64_t>(new_length)) {
+      position_ids_shape_[1] = static_cast<int64_t>(new_length);
+      position_ids_ = OrtValue::CreateTensor(model_.p_device_inputs_->GetAllocator(), position_ids_shape_, position_ids_type_);
+      inputs_[position_ids_index_] = position_ids_.get();
+    }
+
+    const int pos_batch = static_cast<int>(position_ids_shape_[0]);
+    const int pos_seq = static_cast<int>(new_length);
+    if (pos_batch > 1) {
+      // For batch-beam > 1, initialize/update each row via the batch=1 kernel path.
+      // This avoids the generic multi-batch increment path, which is not valid for
+      // prefill matrix initialization in Whisper beam search.
+      auto* base = static_cast<uint8_t*>(position_ids_->GetTensorMutableRawData());
+      const size_t row_bytes = static_cast<size_t>(pos_seq) * Ort::SizeOf(position_ids_type_);
+      for (int row = 0; row < pos_batch; ++row) {
+        void* row_ptr = base + static_cast<size_t>(row) * row_bytes;
+        const bool updated_on_device = model_.p_device_inputs_->UpdatePositionIds(row_ptr,
+                                                                                   1,
+                                                                                   current_length,
+                                                                                   pos_seq,
+                                                                                   position_ids_type_);
+        if (!updated_on_device) {
+          assert(false && "TRT-RTX position_ids row updates must remain on GPU.");
+          throw std::runtime_error("TRT-RTX position_ids row update unexpectedly attempted CPU fallback.");
+        }
+      }
+    } else {
+      const bool updated_on_device = model_.p_device_inputs_->UpdatePositionIds(position_ids_->GetTensorMutableRawData(),
+                                                                                 pos_batch,
+                                                                                 current_length,
+                                                                                 pos_seq,
+                                                                                 position_ids_type_);
+      if (!updated_on_device) {
+        assert(false && "TRT-RTX position_ids updates must remain on GPU.");
+        throw std::runtime_error("TRT-RTX position_ids update unexpectedly attempted CPU fallback.");
+      }
+    }
+  }
+
   // Return early if this method is just initializing the above OrtValue objects and not updating them
   if (first_run_) {
     return;
@@ -131,6 +301,11 @@ void WhisperDecoderState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens,
     *data = current_length - 1;
   }
 
+  // Update attention_mask if model supports it (GQA in-place KV cache)
+  if (has_attention_mask_input_) {
+    UpdateAttentionMask(current_length, static_cast<int>(new_length));
+  }
+
   if (cache_indirection_ && params_->search.num_beams > 1 && !first_update) {
     // Only update after having run one pass through the decoder with past KV caches
     auto beam_indices_span = beam_indices.Span();
@@ -166,6 +341,7 @@ void WhisperDecoderState::UpdateInputsOutputs(DeviceSpan<int32_t>& next_tokens,
       outputs_[output_cross_qk_index_ + i] = output_cross_qk_[i].get();
     }
   }
+
 }
 
 WhisperState::WhisperState(const WhisperModel& model, const GeneratorParams& params, DeviceSpan<int32_t> sequence_lengths_unk)