[WebGPU] Support continuous decoding (RewindTo) with graph capture

src/models/position_inputs.cpp

@@ -380,21 +380,37 @@ void DefaultPositionInputs::InitializeSequenceLengths(std::array<int64_t, 2> sha
 }
 
 void DefaultPositionInputs::RewindMask(size_t index) {
-  if (state_.params_->use_graph_capture) {
-    throw std::runtime_error("PositionInputs::RewindMask - Static buffer is not supported for continuous decoding.");
-#if 0  // TODO: Fix implementation, cudaMemsetAsync of 1 is setting bytes of 1 vs int32's of 1
-    int past_length = static_cast<int>(index);
-    int max_length = static_cast<int>(state_.params_->search.max_length);
-    cudaMemsetAsync(attention_mask_->GetTensorMutableRawData(),
-                    0,
-                    (type_ == Ort::TypeToTensorType<int32_t> ? sizeof(int32_t) : sizeof(int64_t)) * max_length,
-                    model_.cuda_stream_);
-    cudaMemsetAsync(attention_mask_->GetTensorMutableRawData(),
-                    1,
-                    (type_ == Ort::TypeToTensorType<int32_t> ? sizeof(int32_t) : sizeof(int64_t)) * past_length,
-                    model_.cuda_stream_);
-#endif
+  if (ShouldUseStaticMaskHandling()) {
+    // Static mask layout: [batch_beam_size, max_length]
+    // Rewind to index: write 1s for [0, index), 0s for [index, max_length)
+    size_t max_len = static_cast<size_t>(state_.params_->search.max_length);
+    if (index > max_len) {
+      throw std::runtime_error("RewindMask: index exceeds max_length");
+    }
+    size_t batch_beam_size = static_cast<size_t>(attention_mask_shape_[0]);
+    auto byte_span = attention_mask_->GetByteSpan();
+    auto cpu_data = byte_span.CpuSpan();
+    if (type_ == Ort::TypeToTensorType<int32_t>) {
+      auto* data = reinterpret_cast<int32_t*>(cpu_data.data());
+      for (size_t i = 0; i < batch_beam_size; i++) {
+        std::fill_n(data + i * max_len, index, static_cast<int32_t>(1));
+        std::fill_n(data + i * max_len + index, max_len - index, static_cast<int32_t>(0));
+      }
+    } else {
+      auto* data = reinterpret_cast<int64_t*>(cpu_data.data());
+      for (size_t i = 0; i < batch_beam_size; i++) {
+        std::fill_n(data + i * max_len, index, static_cast<int64_t>(1));
+        std::fill_n(data + i * max_len + index, max_len - index, static_cast<int64_t>(0));
+      }
+    }
+    byte_span.CopyCpuToDevice();
+    return;
   }
+
+  // Dynamic mask: adjust shape so the next Update() creates the correct-sized tensor.
+  // For batch_beam_size == 1 (the only case RewindTo supports), the CPU UpdateAttentionMask
+  // fills the entire next mask with 1s, so no data fixup is needed — just the shape.
+  attention_mask_shape_[1] = static_cast<int64_t>(index);
 }
 
 bool DefaultPositionInputs::ShouldUseStaticMaskHandling() const {

src/webgpu/interface.cpp

@@ -171,6 +171,11 @@ struct InterfaceImpl : DeviceInterface {
  private:
   Ort::Allocator* ort_allocator_{};
   const OrtMemoryInfo* ort_memory_info_{};
+  // Reusable CPU staging buffers for UpdateAttentionMask, pre-filled with 1s.
+  // Content is always all 1s so sharing across generators is safe; only upload_bytes
+  // worth of data is copied each call, regardless of buffer capacity.
+  std::vector<int32_t> mask_staging_buffer_i32_;
+  std::vector<int64_t> mask_staging_buffer_i64_;
 
  public:
   Ort::Allocator& GetAllocator() override {
@@ -190,6 +195,44 @@ struct InterfaceImpl : DeviceInterface {
 
   void Synchronize() override {}  // Nothing to do?
 
+  bool UpdateAttentionMask(void* next_mask_data, void* mask_data, int batch_beam_size, int new_kv_length, int total_length, int max_length, bool update_only, ONNXTensorElementDataType type) override {
+    if (batch_beam_size != 1 || !update_only) {
+      return false;  // Fall back to CPU for multi-beam or non-static mask
+    }
+    // For batch_beam_size == 1 with static mask (update_only=true, no padding),
+    // the mask is always all 1s for attended positions.
+    size_t num_elements = static_cast<size_t>(total_length);
+    size_t upload_bytes;
+    void* staging_data;
+
+    // Use the correctly typed staging buffer. Each grows monotonically and
+    // only newly extended positions need to be filled with 1.
+    if (type == Ort::TypeToTensorType<int32_t>) {
+      if (mask_staging_buffer_i32_.size() < num_elements) {
+        mask_staging_buffer_i32_.resize(num_elements, static_cast<int32_t>(1));
+      }
+      staging_data = mask_staging_buffer_i32_.data();
+      upload_bytes = num_elements * sizeof(int32_t);
+    } else {
+      if (mask_staging_buffer_i64_.size() < num_elements) {
+        mask_staging_buffer_i64_.resize(num_elements, static_cast<int64_t>(1));
+      }
+      staging_data = mask_staging_buffer_i64_.data();
+      upload_bytes = num_elements * sizeof(int64_t);
+    }
+
+    int64_t shape_val = static_cast<int64_t>(upload_bytes);
+    std::span<const int64_t> shape{&shape_val, 1};
+    auto cpu_mem_info = OrtMemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);
+    auto src_tensor = OrtValue::CreateTensor(*cpu_mem_info, staging_data, upload_bytes, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8);
+    auto dst_tensor = OrtValue::CreateTensor(*ort_memory_info_, mask_data, upload_bytes, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8);
+    const std::vector<const OrtValue*> src_ptrs = {src_tensor.get()};
+    const std::vector<OrtValue*> dst_ptrs = {dst_tensor.get()};
+    GetOrtEnv().CopyTensors(src_ptrs, dst_ptrs, nullptr);
+
+    return true;
+  }
+
   bool Cast(void* input, void* output, ONNXTensorElementDataType input_type, ONNXTensorElementDataType output_type, size_t element_count) override {
     if (!ort_allocator_) {
       throw std::runtime_error("WebGPU allocator not initialized");

test/c_api_tests.cpp

@@ -1332,8 +1332,197 @@ TEST(CAPITests, RewindGptFp32CAPI) {
   expected_output_start = &expected_output[0];
   EXPECT_TRUE(0 == std::memcmp(expected_output_start, sequence_data, sequence_length * sizeof(int32_t)));
 }
+
+// Test RewindTo(0) with batch=1: full rewind should produce identical output
+TEST(CAPITests, RewindToZeroGptFp32CAPI) {
+  std::vector<int32_t> input_ids{0, 0, 195, 731};
+  std::vector<int32_t> expected_output{0, 0, 195, 731, 731, 114, 114, 114, 114, 114};
+  int max_length = 10;
+
+  auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto params = OgaGeneratorParams::Create(*model);
+  params->SetSearchOption("max_length", max_length);
+
+  auto generator = OgaGenerator::Create(*model, *params);
+  generator->AppendTokens(input_ids.data(), input_ids.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  auto sequence_length = generator->GetSequenceCount(0);
+  auto* sequence_data = generator->GetSequenceData(0);
+  ASSERT_EQ(sequence_length, static_cast<size_t>(max_length));
+  EXPECT_TRUE(0 == std::memcmp(expected_output.data(), sequence_data, sequence_length * sizeof(int32_t)));
+
+  // Full rewind and regenerate — output must be identical
+  generator->RewindTo(0);
+  generator->AppendTokens(input_ids.data(), input_ids.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  sequence_length = generator->GetSequenceCount(0);
+  sequence_data = generator->GetSequenceData(0);
+  ASSERT_EQ(sequence_length, static_cast<size_t>(max_length));
+  EXPECT_TRUE(0 == std::memcmp(expected_output.data(), sequence_data, sequence_length * sizeof(int32_t)));
+}
+
+// Test multiple sequential RewindTo calls: rewind to different positions in succession
+TEST(CAPITests, MultipleRewindGptFp32CAPI) {
+  std::vector<int32_t> input_ids{0, 0, 195, 731};
+  std::vector<int32_t> expected_output{0, 0, 195, 731, 731, 114, 114, 114, 114, 114};
+  int max_length = 10;
+
+  auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto params = OgaGeneratorParams::Create(*model);
+  params->SetSearchOption("max_length", max_length);
+
+  auto generator = OgaGenerator::Create(*model, *params);
+  generator->AppendTokens(input_ids.data(), input_ids.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  // Verify initial generation
+  auto sequence_length = generator->GetSequenceCount(0);
+  ASSERT_EQ(sequence_length, static_cast<size_t>(max_length));
+  EXPECT_TRUE(0 == std::memcmp(expected_output.data(), generator->GetSequenceData(0), sequence_length * sizeof(int32_t)));
+
+  // Rewind to 7, generate remaining
+  generator->RewindTo(7);
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+  sequence_length = generator->GetSequenceCount(0);
+  ASSERT_EQ(sequence_length, static_cast<size_t>(max_length));
+  EXPECT_TRUE(0 == std::memcmp(expected_output.data(), generator->GetSequenceData(0), sequence_length * sizeof(int32_t)));
+
+  // Rewind to 5, generate remaining
+  generator->RewindTo(5);
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+  sequence_length = generator->GetSequenceCount(0);
+  ASSERT_EQ(sequence_length, static_cast<size_t>(max_length));
+  EXPECT_TRUE(0 == std::memcmp(expected_output.data(), generator->GetSequenceData(0), sequence_length * sizeof(int32_t)));
+
+  // Rewind all the way to 0 and regenerate with same input
+  generator->RewindTo(0);
+  generator->AppendTokens(input_ids.data(), input_ids.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+  sequence_length = generator->GetSequenceCount(0);
+  ASSERT_EQ(sequence_length, static_cast<size_t>(max_length));
+  EXPECT_TRUE(0 == std::memcmp(expected_output.data(), generator->GetSequenceData(0), sequence_length * sizeof(int32_t)));
+}
+
+// Test RewindTo with new tokens: rewind to a midpoint and append different tokens
+TEST(CAPITests, RewindAndAppendNewTokensGptFp32CAPI) {
+  std::vector<int32_t> input_ids{0, 0, 195, 731};
+  int max_length = 10;
+
+  auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto params = OgaGeneratorParams::Create(*model);
+  params->SetSearchOption("max_length", max_length);
+
+  auto generator = OgaGenerator::Create(*model, *params);
+  generator->AppendTokens(input_ids.data(), input_ids.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  // Save original sequence
+  auto orig_length = generator->GetSequenceCount(0);
+  std::vector<int32_t> original_sequence(orig_length);
+  std::memcpy(original_sequence.data(), generator->GetSequenceData(0), orig_length * sizeof(int32_t));
+
+  // Rewind to 4 and append different tokens
+  generator->RewindTo(4);
+  std::vector<int32_t> new_tokens{52, 204};
+  generator->AppendTokens(new_tokens.data(), new_tokens.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  // First 4 tokens should match, but subsequent tokens may differ because we changed the input
+  auto new_length = generator->GetSequenceCount(0);
+  ASSERT_EQ(new_length, static_cast<size_t>(max_length));
+  EXPECT_TRUE(0 == std::memcmp(original_sequence.data(), generator->GetSequenceData(0), 4 * sizeof(int32_t)));
+
+  // Rewind to 3 and append the original tokens to verify we get the original output back
+  generator->RewindTo(3);
+  std::vector<int32_t> orig_continuation{731, 731};
+  generator->AppendTokens(orig_continuation.data(), orig_continuation.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  new_length = generator->GetSequenceCount(0);
+  ASSERT_EQ(new_length, static_cast<size_t>(max_length));
+  EXPECT_TRUE(0 == std::memcmp(original_sequence.data(), generator->GetSequenceData(0), new_length * sizeof(int32_t)));
+}
 #endif
 
+// Test RewindTo with static mask handling (graph capture / past-present share buffer).
+// NvTensorRtRtx enables ShouldUseStaticMaskHandling() via past-present share buffer.
+// On main (before our fix), RewindMask throws
+// "Static buffer is not supported for continuous decoding."
+// With our fix, the static mask is properly rewritten with 1s/0s.
+// Skipped when the phi3-fp16-nvtrt model is not available (CI-only model).
+TEST(CAPITests, RewindGraphCaptureNvTensorRtRtxCAPI) {
+  std::string nvtrt_path = MODEL_PATH "hf-internal-testing/phi3-fp16-nvtrt";
+  if (!std::filesystem::exists(nvtrt_path)) {
+    GTEST_SKIP() << "NvTensorRtRtx model not available at " << nvtrt_path;
+  }
+
+  auto config = OgaConfig::Create(nvtrt_path.c_str());
+  config->ClearProviders();
+  config->AppendProvider("NvTensorRtRtx");
+
+  int max_length = 20;
+
+  auto model = OgaModel::Create(*config);
+  auto params = OgaGeneratorParams::Create(*model);
+  params->SetSearchOption("max_length", max_length);
+
+  // Use a simple prompt
+  std::vector<int32_t> input_ids{1, 15043, 29892, 920};  // "Hello, world"
+
+  auto generator = OgaGenerator::Create(*model, *params);
+  generator->AppendTokens(input_ids.data(), input_ids.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  // Save first-run output
+  auto seq_len = generator->GetSequenceCount(0);
+  std::vector<int32_t> first_output(seq_len);
+  std::memcpy(first_output.data(), generator->GetSequenceData(0), seq_len * sizeof(int32_t));
+
+  // RewindTo(0) — full rewind with static mask. This threw on main.
+  generator->RewindTo(0);
+  generator->AppendTokens(input_ids.data(), input_ids.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  // Output after rewind must match first run
+  auto seq_len2 = generator->GetSequenceCount(0);
+  ASSERT_EQ(seq_len2, seq_len);
+  EXPECT_TRUE(0 == std::memcmp(first_output.data(), generator->GetSequenceData(0), seq_len * sizeof(int32_t)));
+
+  // RewindTo(6) — partial rewind with static mask
+  generator->RewindTo(6);
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  seq_len2 = generator->GetSequenceCount(0);
+  ASSERT_EQ(seq_len2, seq_len);
+  EXPECT_TRUE(0 == std::memcmp(first_output.data(), generator->GetSequenceData(0), seq_len * sizeof(int32_t)));
+}
+
 #ifndef STREAMING_ASR_PATH
 #define STREAMING_ASR_PATH MODEL_PATH "nemotron-speech-streaming"
 #endif