Harden Engine request lifetime and sampler acquisition

docs/paged_attention_engine.md

@@ -12,7 +12,7 @@ The current dynamic path manages paged KV decoder state together with per-reques
 
 > **Transitional low-level API:** `AddTokens()` plus `AddRequest()`, `Continue()`, repeated `Step()` calls, token-at-a-time unseen-output access, and `Remove()` are a transitional host-facing surface. The production host API is expected to wrap or replace these operations; do not treat their current shape as the final high-level contract.
 >
-> **Serialization requirement:** Except for releasing an external request handle, every call on an `Engine` and on any `Request` owned by that engine must be externally serialized with `Engine::Step()`. This includes completion and unseen-output access as well as lifecycle mutation. Final handle release only publishes an atomic abandonment marker; cleanup runs at the next serialized Engine boundary. The API is otherwise not thread-safe, and idempotent terminal removal only makes sequential retries harmless.
+> **Serialization requirement:** Except for releasing an external request handle, every call on an `Engine` and on any `Request` owned by that engine must be externally serialized with `Engine::Step()`. This includes completion and unseen-output access as well as lifecycle mutation. Final handle release only publishes an abandonment marker; cleanup runs at the next serialized Engine boundary. External handle zero/one transitions serialize the self-owner and base-owned lifecycle state, so a concurrent handle returned by `Engine::Step()` cannot be erased by the previous handle's final release. The API is otherwise not thread-safe, and idempotent terminal removal only makes sequential retries harmless.
 
 The main implementation is under `src/engine/`:
 
@@ -63,6 +63,13 @@ return ready requests one at a time
 
 The step is transactional. Planning and reservation do not immediately change committed request or cache state. If a recoverable failure occurs before commit, the engine restores the request search state and releases the reserved cache blocks. A failure during the commit boundary is considered fatal because the engine can no longer guarantee that all cooperating components agree on the committed state.
 
+Diagnostic invariant snapshots cross-check each committed full-cache table's used slots against its
+Request's `processed_sequence_length_`. For windowed caches, the full and ring-cache owner sets must
+match. Full-cache tables, ring-cache tables, and active reservation deltas must all refer to known
+Request snapshots, and each reservation records both its full-cache and window-cache block ownership
+so inconsistent membership or unattributed reserved blocks are detectable. These checks are test and
+diagnostic machinery; they do not add validation to the runtime hot path.
+
 ## How the dynamic path is selected
 
 `Engine::CreateDependencies()` creates three collaborators from the model:
@@ -132,15 +139,20 @@ The request is not owned by an engine. `AddTokens()` accumulates the initial pro
 
 ### `Assigned`
 
-`Engine::AddRequest()` validates the request, calls `Request::Assign()`, and adds it to the scheduler pool.
+`Engine::AddRequest()` validates the request, prepares a detached search sequence, device prompt,
+host mirrors, sampler state, scheduler capacity, and Engine tracking capacity, then commits the
+request and inserts it into the scheduler using nonthrowing moves into reserved storage.
 
 This is the queued state. `Engine::AddRequest()` moves a new request here before first admission.
 `Continue()` also moves a cache-resident `TurnComplete` request here while its next input waits for
 execution.
 
-For a new request, assignment moves the prompt into `Search`, creates the host-side token mirror,
-initializes the sequence counters, and records the owning Engine. `AddTokens()` and `Continue()` are
-both rejected while already queued. Input must leave room for at least one generated token below
+For a new request, admission does not publish Engine ownership or queued status until all request and
+scheduler preparation succeeds. A preparation failure therefore leaves the request `Unassigned`
+with its original prompt intact and eligible for retry. Commit moves the prepared prompt into
+`Search`, installs the host-side token mirror and sampler state, initializes the sequence counters,
+records the owning Engine, and inserts the request into the scheduler. `AddTokens()` and `Continue()`
+are both rejected while already queued. Input must leave room for at least one generated token below
 `max_length`.
 
 `max_length` is the cumulative total sequence limit for the entire session: the initial prompt, generated output, and every continuation input all count against the same limit. `Continue()` does not reset it, and it is not a per-turn generation budget.
@@ -164,6 +176,10 @@ the model's chat template.
 
 A submitted request remains owned by its Engine and resident at `TurnComplete`. `Remove()` releases that ownership immediately. If every external handle is released instead, the request is marked abandoned and reclaimed before the Engine's next `AddRequest()` or `Step()` boundary.
 
+Automatic abandonment on final handle destruction is transitional behavior, not logical close. The
+external-reference machinery keeps that behavior race-free today, but the production ownership layer
+should keep handle destruction separate from explicit request close rather than reusing this policy.
+
 Planning skips turn-complete residents and does not release their cache. Retained requests still
 consume paged-cache blocks and a batch slot, so applications must call `Remove()` when they no
 longer need continuation when deterministic immediate reclamation is required.

src/cuda/interface.cpp

@@ -8,12 +8,17 @@
 #include "search_cuda.h"
 #include "kernels.h"
 #include "cuda_topk.h"
+#include <algorithm>
+#include <cassert>
 #include <charconv>
+#include <cstddef>
 #include <cstdarg>
 #include <cstring>
 #include <mutex>
 #include <random>
 #include <system_error>
+#include <utility>
+#include <vector>
 
 #if defined(_WIN32) || defined(_WIN64)
 #define strcasecmp _stricmp
@@ -89,6 +94,52 @@ DeviceSpan<T> AllocateCudaSpan(size_t count) {
   return DeviceSpan<T>{std::make_shared<GpuMemory>(count * sizeof(T))};
 }
 
+namespace {
+
+class SamplerStateIndexPool {
+ public:
+  template <typename Prepare, typename Create>
+  auto AcquireOwned(Prepare&& prepare, Create&& create) {
+    const bool reusing = !free_indices_.empty();
+    const int index = reusing ? free_indices_.back() : size_;
+    const int required_size = reusing ? size_ : size_ + 1;
+
+    // Release must remain allocation-free, so reserve its future slot before any external
+    // preparation can publish an acquired index.
+    free_indices_.reserve(static_cast<size_t>(required_size));
+    std::forward<Prepare>(prepare)(index, required_size);
+
+    if (reusing) {
+      free_indices_.pop_back();
+    } else {
+      size_ = required_size;
+    }
+
+    try {
+      return std::forward<Create>(create)(index);
+    } catch (...) {
+      Release(index);
+      throw;
+    }
+  }
+
+  void Release(int index) noexcept {
+    assert(index >= 0 && index < size_);
+    assert(std::find(free_indices_.begin(), free_indices_.end(), index) ==
+           free_indices_.end());
+    assert(free_indices_.size() < free_indices_.capacity());
+    free_indices_.push_back(index);
+  }
+
+  int Size() const noexcept { return size_; }
+
+ private:
+  std::vector<int> free_indices_;
+  int size_{};
+};
+
+}  // namespace
+
 struct CudaSamplerStatePool {
   explicit CudaSamplerStatePool(int initial_capacity) {
     if (initial_capacity > 0) {
@@ -97,26 +148,22 @@ struct CudaSamplerStatePool {
     }
   }
 
-  int Acquire(int random_seed) {
-    int index;
-    if (free_indices_.empty()) {
-      index = size_++;
-      EnsureCapacity(size_);
-    } else {
-      index = free_indices_.back();
-      free_indices_.pop_back();
-    }
-
-    const unsigned long long seed = random_seed == -1
-                                        ? static_cast<unsigned long long>(std::random_device{}())
-                                        : static_cast<unsigned long long>(random_seed);
-    cuda::LaunchInitCurandState(seed, states_.Span().data() + index, GetStream());
-    return index;
+  template <typename Create>
+  auto AcquireOwned(int random_seed, Create&& create) {
+    return indices_.AcquireOwned(
+        [this, random_seed](int index, int required_size) {
+          EnsureCapacity(required_size);
+          const unsigned long long seed =
+              random_seed == -1
+                  ? static_cast<unsigned long long>(std::random_device{}())
+                  : static_cast<unsigned long long>(random_seed);
+          cuda::LaunchInitCurandState(
+              seed, states_.Span().data() + index, GetStream());
+        },
+        std::forward<Create>(create));
   }
 
-  void Release(int index) {
-    free_indices_.push_back(index);
-  }
+  void Release(int index) noexcept { indices_.Release(index); }
 
   curandState* Data() { return states_.Span().data(); }
 
@@ -127,9 +174,9 @@ struct CudaSamplerStatePool {
 
     const int new_capacity = std::max(required_capacity, std::max(4, capacity_ * 2));
     auto new_states = AllocateCudaSpan<curandState>(new_capacity);
-    if (size_ > 1) {
+    if (indices_.Size() > 0) {
       CUDA_CHECK(cudaMemcpyAsync(new_states.Span().data(), states_.Span().data(),
-                                 static_cast<size_t>(size_ - 1) * sizeof(curandState),
+                                 static_cast<size_t>(indices_.Size()) * sizeof(curandState),
                                  cudaMemcpyDeviceToDevice, GetStream()));
       CUDA_CHECK(cudaStreamSynchronize(GetStream()));
     }
@@ -138,16 +185,15 @@ struct CudaSamplerStatePool {
   }
 
   DeviceSpan<curandState> states_;
-  std::vector<int> free_indices_;
-  int size_{};
+  SamplerStateIndexPool indices_;
   int capacity_{};
 };
 
 struct CudaBatchedSamplerState final : BatchedSamplerState {
   CudaBatchedSamplerState(std::shared_ptr<CudaSamplerStatePool> pool, int index)
       : pool_{std::move(pool)}, index_{index} {}
 
-  ~CudaBatchedSamplerState() override { pool_->Release(index_); }
+  ~CudaBatchedSamplerState() noexcept override { pool_->Release(index_); }
 
   std::shared_ptr<CudaSamplerStatePool> pool_;
   int index_{};
@@ -169,7 +215,12 @@ struct CudaBatchedSampler final : BatchedSampler {
   }
 
   std::unique_ptr<BatchedSamplerState> CreateState(int random_seed) override {
-    return std::make_unique<CudaBatchedSamplerState>(state_pool_, state_pool_->Acquire(random_seed));
+    auto pool = state_pool_;
+    return state_pool_->AcquireOwned(
+        random_seed,
+        [pool = std::move(pool)](int index) {
+          return std::make_unique<CudaBatchedSamplerState>(pool, index);
+        });
   }
 
   bool OwnsState(const BatchedSamplerState& state) const override {

src/engine/engine.cpp

@@ -60,24 +60,22 @@ void Engine::AddRequest(std::shared_ptr<Request> request) {
     request->ValidateEngineCompatibility();
   }
 
-  // Track the request before assignment so every successfully submitted request can later be found
-  // even when the scheduler and cache hold it through implementation-specific containers. The
-  // registry allocation therefore happens before any request lifecycle mutation.
-  tracked_requests_.push_back(request);
-  try {
-    request->Assign(shared_from_this());
-    scheduler_->AddRequest(request);
-  } catch (...) {
-    tracked_requests_.pop_back();
-    throw;
-  }
+  auto request_preparation = request->PrepareAdmission();
+  tracked_requests_.reserve(tracked_requests_.size() + 1);
+  auto scheduler_preparation = scheduler_->PrepareAddRequest(request);
+
+  request_preparation.sampling_state =
+      std::move(scheduler_preparation.sampling_state);
+  request->CommitAdmission(shared_from_this(), std::move(request_preparation));
+  scheduler_->CommitAddRequest(request, std::move(scheduler_preparation));
+  tracked_requests_.emplace_back(request);
 }
 
 void Engine::RemoveRequest(std::shared_ptr<Request> request) {
   if (request && IsClosed(request->status_)) {
     return;
   }
-  if (!request || request->engine_.lock().get() != this) {
+  if (!request || !request->BelongsTo(*this)) {
     throw std::runtime_error("Cannot remove a request from an engine it does not belong to.");
   }
 
@@ -93,7 +91,7 @@ void Engine::RemoveRequest(std::shared_ptr<Request> request) {
   staged_ready_requests_.erase(
       std::remove(staged_ready_requests_.begin(), staged_ready_requests_.end(), request),
       staged_ready_requests_.end());
-  request->CompleteClose();
+  request->CompleteCloseFromEngine(*this);
   tracked_requests_.erase(
       std::remove_if(
           tracked_requests_.begin(), tracked_requests_.end(),
@@ -105,8 +103,8 @@ void Engine::RemoveRequest(std::shared_ptr<Request> request) {
 }
 
 void Engine::ReclaimAbandonedRequests() {
-  // ExternalRelease only publishes an atomic abandonment marker. Engine entry points are externally
-  // serialized, so this boundary can safely perform the normal removal sequence: scheduler/cache
+  // ExternalRelease only publishes synchronized external-lifecycle state. Engine entry points are
+  // externally serialized, so this boundary can safely perform the normal removal sequence: scheduler/cache
   // release, ready-notification purge, and terminal close.
   std::vector<std::shared_ptr<Request>> abandoned_requests;
   abandoned_requests.reserve(tracked_requests_.size());
@@ -119,8 +117,8 @@ void Engine::ReclaimAbandonedRequests() {
               return true;
             }
             if (!IsClosed(request->status_) &&
-                request->engine_.lock().get() == this &&
-                request->IsExternallyAbandoned()) {
+                request->BelongsTo(*this) &&
+                request->ExternalReferencesAbandoned()) {
               abandoned_requests.push_back(request);
             }
             return false;
@@ -129,7 +127,7 @@ void Engine::ReclaimAbandonedRequests() {
 
   for (const auto& request : abandoned_requests) {
     // Recheck defensively in case an external owner was reacquired before this serialized boundary.
-    if (request->IsExternallyAbandoned()) {
+    if (request->ExternalReferencesAbandoned()) {
       RemoveRequest(request);
     }
   }
@@ -139,7 +137,7 @@ void Engine::ValidateRequestCanContinue(const std::shared_ptr<Request>& request)
   if (health_ == EngineHealth::Unhealthy) {
     std::rethrow_exception(fatal_error_);
   }
-  if (request->engine_.lock().get() != this) {
+  if (!request->BelongsTo(*this)) {
     throw std::runtime_error("Cannot continue a request that does not belong to this engine.");
   }
 
@@ -174,7 +172,7 @@ void Engine::ValidateRequestCanContinue(const std::shared_ptr<Request>& request)
     message = AddExceptionCause(
         std::move(message) + " Closing the poisoned request also failed.",
         std::current_exception());
-    request->CompleteClose();
+    request->CompleteCloseFromEngine(*this);
   }
   MarkUnhealthyAndThrow(
       StepOutcomeKind::FatalExecutionFailure,

src/engine/engine.h

@@ -118,16 +118,19 @@ struct Engine : std::enable_shared_from_this<Engine>,
    */
   bool HasPendingRequests() const;
 
- private:
-  void ReclaimAbandonedRequests();
-  std::shared_ptr<Request> DrainReadyRequest();
-  std::shared_ptr<Request> StepDynamic();
-  std::shared_ptr<Request> StepStatic();
+  // Internal continuation preflight used by Request::Continue(). It validates ownership, health,
+  // residency, ready-drain ordering, and static-batch constraints without exposing scheduler state.
   void ValidateRequestCanContinue(const std::shared_ptr<Request>& request) const;
   [[noreturn]] void HandleContinuationRestoreFailure(
       const std::shared_ptr<Request>& request,
       std::exception_ptr append_error,
       std::exception_ptr restore_error);
+
+ private:
+  void ReclaimAbandonedRequests();
+  std::shared_ptr<Request> DrainReadyRequest();
+  std::shared_ptr<Request> StepDynamic();
+  std::shared_ptr<Request> StepStatic();
   [[noreturn]] void MarkUnhealthyAndThrow(StepOutcomeKind outcome,
                                           StepTransactionId transaction_id,
                                           const void* request_id,
@@ -148,8 +151,6 @@ struct Engine : std::enable_shared_from_this<Engine>,
   std::vector<std::shared_ptr<Request>> ready_requests_;
   std::vector<std::shared_ptr<Request>> staged_ready_requests_;
   size_t ready_request_index_{};
-
-  friend struct Request;
 };
 
 }  // namespace Generators

src/engine/paged_key_value_cache.cpp

@@ -20,12 +20,6 @@ StateGroup ResolvePagedKeyValueGroup(const Config::Model::Decoder& decoder) {
   if (!decoder.state_groups) {
     StateGroup group;
     group.kind = StateGroupKind::PagedKeyValue;
-    group.key = Config::Model::Decoder::StateBinding{
-        decoder.inputs.past_key_names,
-        decoder.outputs.present_key_names};
-    group.value = Config::Model::Decoder::StateBinding{
-        decoder.inputs.past_value_names,
-        decoder.outputs.present_value_names};
     group.layer_ids.reserve(decoder.num_hidden_layers);
     for (int layer_id = 0; layer_id < decoder.num_hidden_layers; ++layer_id) {
       group.layer_ids.push_back(layer_id);
@@ -52,19 +46,18 @@ StateGroup ResolvePagedKeyValueGroup(const Config::Model::Decoder& decoder) {
     throw std::runtime_error(
         "Dynamic batching requires one paged_kv decoder state group");
   }
-  if (!paged_group->key || !paged_group->value ||
-      paged_group->layer_ids.empty()) {
+  if (paged_group->layer_ids.empty()) {
     throw std::runtime_error(
-        "Dynamic batching requires a non-empty paged_kv decoder state group "
-        "with key and value bindings");
+        "Dynamic batching requires a non-empty paged_kv decoder state group");
   }
   return *paged_group;
 }
 
 ONNXTensorElementDataType KeyValueCacheType(const std::shared_ptr<Model>& model,
                                             const StateGroup& paged_group) {
   const auto key_name = ComposeKeyValueName(
-      paged_group.key->input, paged_group.layer_ids.front());
+      model->config_->model.decoder.inputs.past_key_names,
+      paged_group.layer_ids.front());
   return model->session_info_.GetInputDataType(key_name);
 }
 
@@ -227,10 +220,10 @@ PagedKeyValueCache::PagedKeyValueCache(std::shared_ptr<Model> model)
     cache_.push_back(LayerCache{
         OrtValue::CreateTensor(model->p_device_kvcache_->GetAllocator(), cache_shape_per_layer, dtype),  // Key cache
         OrtValue::CreateTensor(model->p_device_kvcache_->GetAllocator(), cache_shape_per_layer, dtype),  // Value cache
-        ComposeKeyValueName(paged_group.key->input, layer_id),
-        ComposeKeyValueName(paged_group.value->input, layer_id),
-        ComposeKeyValueName(paged_group.key->output, layer_id),
-        ComposeKeyValueName(paged_group.value->output, layer_id)});
+        ComposeKeyValueName(decoder.inputs.past_key_names, layer_id),
+        ComposeKeyValueName(decoder.inputs.past_value_names, layer_id),
+        ComposeKeyValueName(decoder.outputs.present_key_names, layer_id),
+        ComposeKeyValueName(decoder.outputs.present_value_names, layer_id)});
   }
   block_pool_ = std::make_unique<BlockPool>(block_size, num_blocks);
   if (Windowed()) {