From bd15d40d559ed4b365c28366d7e99648d91abf4a Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 01:04:04 -0500 Subject: [PATCH 01/19] Add Engine continuous decoding lifecycle Separate initial input from multi-turn continuation so Engine requests can retain and reuse model state safely across turns. - add Created/Queued/InProgress/TurnComplete/Closed lifecycle and Continue APIs across C, C++, and Python - preserve resident cache state, unseen output ordering, scheduler transactions, and explicit Remove semantics - cover static and dynamic scheduling, rollback, backpressure, ready queues, docs, examples, and integration tests Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b232f82c-f25c-429a-a855-7f8c8f40bbf3 --- docs/paged_attention_batching.md | 6 +- docs/paged_attention_engine.md | 120 +++++-- examples/python/engine/model-qa.py | 81 +++-- src/engine/cache_manager.cpp | 18 +- src/engine/cache_manager.h | 12 + src/engine/decoders/simple_decoder.cpp | 2 +- .../decoders/static_batch_decoder_io.cpp | 21 +- src/engine/decoders/varlen_decoder_io.cpp | 10 +- src/engine/engine.cpp | 81 ++++- src/engine/engine.h | 3 + src/engine/engine_invariants.h | 2 +- src/engine/request.cpp | 130 ++++++-- src/engine/request.h | 44 ++- src/engine/request_status.h | 30 +- src/engine/scheduled_requests.cpp | 23 +- src/engine/scheduler.cpp | 94 +++--- src/engine/scheduler.h | 7 +- src/generators.cpp | 27 +- src/generators.h | 1 + src/ort_genai.h | 10 + src/ort_genai_c.cpp | 37 ++- src/ort_genai_c.h | 46 ++- src/python/python.cpp | 14 + test/engine/engine_invariants_tests.cpp | 6 +- test/engine/engine_step_tests.cpp | 305 +++++++++++++++++- test/engine/engine_test_doubles.h | 6 + test/engine/request_lifecycle_tests.cpp | 150 ++++++++- test/engine/scheduler_contract_tests.cpp | 82 ++++- test/python/test_onnxruntime_genai_engine.py | 116 +++++++ 29 files changed, 1263 insertions(+), 221 deletions(-) diff --git a/docs/paged_attention_batching.md b/docs/paged_attention_batching.md index 6338c68265..64ffd736e7 100644 --- a/docs/paged_attention_batching.md +++ b/docs/paged_attention_batching.md @@ -238,8 +238,8 @@ inheriting any of the constraints that make it unusable for a server. The fast path is taken only when all of the following hold for the current step. Otherwise the existing two-phase per-request loop runs unchanged. -- At least two scheduled requests, none of them already `Completed`. A completed request is skipped - by the per-request loops, which would leave a hole in the logits rows that a single batched +- At least two executable scheduled requests. `TurnComplete` and `Closed` rows are skipped by the + per-request loops, which would otherwise leave holes in the logits rows that a single batched sampler call cannot express. - Every scheduled request resolves to the same `(k, p, temperature)` triple. `Request` funnels every sampling branch into `SampleTopKTopP`, so comparing the resolved triple rather than the raw @@ -392,7 +392,7 @@ Generator design" option in its strongest form. It was rejected because it requi `Sequences` to carry a per-row length cursor (today: one `current_length_`, and `GetSequence(i)` depends on it), reworking every kernel that writes at a shared `past_length` offset, moving per-request `GeneratorParams` into per-row arrays, and adding row allocation/eviction to `Search`. -It also collides with `Request`'s public lifecycle β€” `Assign`, `Remove`, `AddTokens` can all be +It also collides with `Request`'s public lifecycle β€” `Assign`, `Remove`, `AddTokens`, and `Continue` can all be called outside the engine. It is a plausible long-term direction, but it is a rewrite of the search layer, and Phase 2 gets most of the benefit without touching any of it. diff --git a/docs/paged_attention_engine.md b/docs/paged_attention_engine.md index 9f6f8a163b..ac08d75bcf 100644 --- a/docs/paged_attention_engine.md +++ b/docs/paged_attention_engine.md @@ -92,33 +92,74 @@ The engine creates throughput by batching several independent requests, not by p The important request states are: ```text -Unassigned -> Assigned -> InProgress -> Completed - ^ | | | - +-----------+------------+------------+ - Remove() +Unassigned (Created) -- submit --> Assigned (Queued) -- schedule --> InProgress + ^ | + | | turn stops + +---- Continue(tokens) ---- TurnComplete + +Assigned (Queued) ---+ +InProgress ----------+-- Remove() --> Closed +TurnComplete --------+ ``` ### `Unassigned` -The request is not owned by an engine. Input tokens added in this state are kept as prefill input. +The request is not owned by an engine. `AddTokens()` accumulates the initial prompt in this state. +`Continue()` is not valid until a submitted request reaches `TurnComplete`. ### `Assigned` `Engine::AddRequest()` validates the request, calls `Request::Assign()`, and adds it to the scheduler pool. -Assignment moves the prompt into the request's `Search`, creates the host-side token mirror, initializes the sequence counters, and records the owning engine. The request has not yet been admitted to the paged cache. +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 +`max_length`. ### `InProgress` -The request has completed at least one committed engine transaction and belongs to the active engine workload. It normally has one unprocessed token at the beginning of a decode step: the token sampled by the previous step. +The current turn is executable and owned by the Engine. It normally has one unprocessed token at +the beginning of a decode step: the token sampled by the previous step. + +### `TurnComplete` + +The current generation turn reached an end condition, such as EOS or maximum length. Generated +output remains available, and `IsDone()` means this state rather than permanent request termination. + +A generated EOS/stop token is not appended to the logical sequence or returned as unseen output. +The next continuation fragment is therefore responsible for any turn-boundary tokens required by +the model's chat template. + +`Continue(tokens)` appends the next input fragment and moves a resident request back to `Assigned`. +`AddTokens()` remains an initial-input-only operation. + +There is no fixed wall-clock or next-step timeout. 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. Until Phase 2 defines +residency and eviction, insufficient capacity is surfaced as backpressure rather than silently +discarding another conversation's model state. + +### `Remove()` -### `Completed` +`Remove()` is legal from `Assigned`, `InProgress`, and `TurnComplete`, and moves the request to +terminal `Closed`. On the dynamic path, removal immediately erases scheduler membership and releases +committed paged-cache ownership. The Engine also removes any undrained ready-queue entries for that +request. -The search has reached an end condition, such as an EOS token or maximum length. The request can be returned to the caller immediately, but its cache blocks are normally reclaimed by `DynamicBatchScheduler::ReapCompletedRequests()` at the beginning of the next planning pass. +### `Closed` -`Remove()` is legal from `Assigned`, `InProgress`, and `Completed`, and returns the request to `Unassigned`. On the dynamic path, removal immediately erases scheduler membership and releases committed paged-cache ownership. +`Closed` is distinct from `Unassigned` because removal may already have destroyed residency and +scheduler ownership. Returning to `Unassigned` would imply that the same logical sequence could be +submitted as a new request. A closed static-batch row may remain physically allocated until the +batch is recycled, but it is no longer sampled or returned. -This removal does not purge an entry already placed in `Engine::ready_requests_`. Because `Engine::Step()` drains that queue before scheduling new work, a removed request that was already ready can still be returned by a later `Step()` call. +Lifecycle status and residency are separate concepts. Phase 1 guarantees that dynamic +`TurnComplete` requests stay resident until `Remove()`. Phase 2 will define observable residency and +automatic eviction; no eviction policy is part of this lifecycle change. ## The request length counters @@ -128,7 +169,7 @@ Three views of request progress are important: | --- | --- | | `CurrentSequenceLength()` | Number of tokens currently held by the request's search sequence | | `processed_sequence_length_` | Number of sequence tokens already represented in the committed KV cache | -| `seen_sequence_length_` | Number of tokens already observed by the API caller | +| `seen_sequence_length_` | High-water sequence index of generated output consumed by the API caller; continuation input may create gaps | The unprocessed tokens are: @@ -172,7 +213,7 @@ This distinction is important: - One call to `Step()` does not always mean one model invocation. - Draining a previously committed batch does not change model or cache state. -- `Engine::RemoveRequest()` does not purge requests already in the ready queue; `Step()` drains those entries before scheduling new work. +- `Engine::RemoveRequest()` purges undrained entries for that request from the ready queue. - `HasPendingRequests()` is true while either the ready queue or scheduler contains work. If the engine has previously encountered a fatal transaction or execution failure, `Step()` rethrows the stored error instead of attempting more work. @@ -181,17 +222,24 @@ If the engine has previously encountered a fatal transaction or execution failur `Engine::StepDynamic()` coordinates the complete transaction. -### 1. Reap completed requests +### 1. Identify active and waiting requests -`DynamicBatchScheduler::PlanStep()` begins by calling `ReapCompletedRequests()`. +`DynamicBatchScheduler::PlanStep()` skips `TurnComplete` residents and builds candidates from +executable residents plus waiting requests. -Completed requests that still own paged-cache blocks are deallocated and removed from the scheduler pool. Their released blocks are immediately available when the same planning pass considers new requests. +The cache manager checks whether those candidates fit alongside dormant turn-complete requests. It +does not reclaim another request as a side effect of `Step()`. If retained residency prevents +admission or cache growth, the plan reports capacity backpressure; the application decides which +conversation to release with `Remove()`. ### 2. Build the initial step plan -The scheduler snapshots all requests that already belong to the paged cache. These requests are expected to be `InProgress`. +The scheduler snapshots requests that already belong to the paged cache. Executable residents may +be `InProgress` or `Assigned`; an `Assigned` resident is a queued continuation. -It then snapshots waiting requests from the scheduler pool. These requests are expected to be `Assigned` and are marked as newly admitted candidates. +It then snapshots nonresident waiting requests from the scheduler pool. These are `Assigned` and +are marked as newly admitted candidates. Residency, not status alone, determines +`newly_admitted`. The scheduler orders candidates with decodes first. Order remains stable among decodes and among prefills. Each candidate initially contributes one provisional @@ -383,9 +431,11 @@ Committing request bookkeeping: - Appends the staged token to the host token mirror. - Sets `processed_sequence_length_` to the sequence length that existed before sampling. -- Changes the status to `InProgress` or `Completed`. +- Changes the status to `InProgress` or `TurnComplete`. -For a newly admitted request, this commit is the point where it moves directly from `Assigned` to `InProgress` or `Completed`. The dynamic transaction path does not need a separate visible scheduling state between those two states. +For a new request or queued continuation, this commit is the point where it moves from `Assigned` +to `InProgress` or `TurnComplete`. The dynamic transaction path does not need a separate visible +scheduling state between those states. Finally, the engine swaps the staged ready list into `ready_requests_`. The first ready request is returned immediately, and later calls drain the rest without another model run. @@ -567,6 +617,11 @@ Requests skipped because of token, row, or temporary cache capacity remain pendi If no request can run because of temporary capacity, `StepDynamic()` reports `CapacityDeferred` instead of returning `nullptr`. Returning `nullptr` would incorrectly tell the caller that no work remains. +The native Engine exposes `CapacityDeferred` as a structured `StepOutcomeKind`. The current C and +Python wrappers still surface it as an error message. Phase 2 must add a structured public +backpressure/residency signal before introducing automatic eviction, so applications can select a +turn-complete request to close without parsing text. + ## Static engine path The static engine path is intentionally separate. @@ -575,6 +630,15 @@ The static engine path is intentionally separate. `StepStatic()` performs decode and sampling directly without the dynamic transaction and reservation protocol. +A resident static request queued by `Continue()` returns to `InProgress` without reallocating the +batch. Static cache rows still cannot be released independently, and an all-turn-complete batch may +be recycled for new work. Static continuation is therefore valid only while the original +single-request batch remains resident. + +A closed static row remains physically retained until that shared batch is recycled. It is not +sampled or returned again, but its Request/Search storage can remain alive for the lifetime of the +batch. + Changes to shared types such as `Request`, `ScheduledRequests`, `ModelExecutor`, or `SimpleDecoder` should be checked against both paths. This document should be updated only where behavior is shared or where the dynamic path changes. ## Public API shape @@ -582,6 +646,7 @@ Changes to shared types such as `Request`, `ScheduledRequests`, `ModelExecutor`, The language bindings expose the same basic loop: ```python +request.add_tokens(initial_tokens) engine.add_request(request) while engine.has_pending_requests(): @@ -590,9 +655,22 @@ while engine.has_pending_requests(): while ready_request.has_unseen_tokens(): token = ready_request.get_unseen_token() # Stream or process the token. + +if request.status == og.RequestStatus.TURN_COMPLETE: + request.continue_with(next_turn_tokens) + +# Repeat engine.step(), then close the conversation when continuation is no longer needed. +engine.remove_request(request) ``` -One ready request may be returned several times over its lifetime as new tokens become available. The request remains owned by the engine until it completes or is explicitly removed. +`AddTokens` is for initial input. The explicit continuation operations are `OgaRequestContinue` in +C, `OgaRequest::Continue` in the C++ wrapper, and `request.continue_with` in Python. Lifecycle is +available through `OgaRequestGetStatus`, `OgaRequest::GetStatus`, and `request.status`. +`IsDone()` remains a compatibility convenience for β€œthe current turn is complete.” + +One ready request may be returned several times over its lifetime as new tokens become available. A +turn-complete dynamic request remains cache-resident until explicit removal, which releases dynamic +cache ownership immediately. ## Keeping this document current diff --git a/examples/python/engine/model-qa.py b/examples/python/engine/model-qa.py index fc4105a6b8..1e02ffd19f 100644 --- a/examples/python/engine/model-qa.py +++ b/examples/python/engine/model-qa.py @@ -17,42 +17,53 @@ def run(args: argparse.Namespace): tokenizer = og.Tokenizer(model) engine = og.Engine(model) - while prompt := input("🫡 : "): - if prompt == "/exit": - break - - messages = [ - {"role": "system", "content": ""}, - {"role": "user", "content": f"{prompt}"}, - ] - messages = json.dumps(messages) - - params = og.GeneratorParams(model) - params.set_search_options( - do_sample=False, - max_length=1024, - ) - - request = og.Request(params) - request.add_tokens( - tokenizer.encode(tokenizer.apply_chat_template(messages=messages, add_generation_prompt=True)), - ) - streaming_tokenizer = tokenizer.create_stream() - - engine.add_request(request) - - print("πŸ€– :", end="", flush=True) - - while ready_request := engine.step(): - while ready_request.has_unseen_tokens(): - print( - streaming_tokenizer.decode(ready_request.get_unseen_token()), - end="", - flush=True, - ) + params = og.GeneratorParams(model) + params.set_search_options( + do_sample=False, + max_length=1024, + ) - print() - engine.remove_request(request) + request = og.Request(params) + system_message = json.dumps([{"role": "system", "content": ""}]) + request.add_tokens( + tokenizer.encode( + tokenizer.apply_chat_template(messages=system_message, add_generation_prompt=False), + ), + ) + streaming_tokenizer = tokenizer.create_stream() + request_added = False + + try: + while prompt := input("🫡 : "): + if prompt == "/exit": + break + + user_message = json.dumps([{"role": "user", "content": prompt}]) + turn_tokens = tokenizer.encode( + tokenizer.apply_chat_template(messages=user_message, add_generation_prompt=True), + ) + + if request_added: + request.continue_with(turn_tokens) + else: + request.add_tokens(turn_tokens) + engine.add_request(request) + request_added = True + + print("πŸ€– :", end="", flush=True) + + while ready_request := engine.step(): + while ready_request.has_unseen_tokens(): + print( + streaming_tokenizer.decode(ready_request.get_unseen_token()), + end="", + flush=True, + ) + + print() + finally: + if request_added: + engine.remove_request(request) if __name__ == "__main__": diff --git a/src/engine/cache_manager.cpp b/src/engine/cache_manager.cpp index bdfe8a13cf..bf240d2243 100644 --- a/src/engine/cache_manager.cpp +++ b/src/engine/cache_manager.cpp @@ -83,7 +83,8 @@ bool StaticCacheManager::CanAllocate(const std::vector> if (std::all_of(cache_allocated_requests_.begin(), cache_allocated_requests_.end(), [](const std::shared_ptr& request) { - return request->status_ == RequestStatus::Completed; + return IsTurnComplete(request->status_) || + IsClosed(request->status_); })) { return true; } @@ -97,7 +98,8 @@ void StaticCacheManager::Allocate(const std::vector>& r if (!cache_allocated_requests_.empty() && std::all_of(cache_allocated_requests_.begin(), cache_allocated_requests_.end(), [](const std::shared_ptr& request) { - return request->status_ == RequestStatus::Completed; + return IsTurnComplete(request->status_) || + IsClosed(request->status_); })) { // If all requests are completed, we can deallocate them before allocating the new requests. Deallocate(cache_allocated_requests_); @@ -157,6 +159,11 @@ std::vector> StaticCacheManager::AllocatedRequests() co return cache_allocated_requests_; } +bool StaticCacheManager::IsResident(const std::shared_ptr& request) const { + return std::find(cache_allocated_requests_.begin(), cache_allocated_requests_.end(), request) != + cache_allocated_requests_.end(); +} + PagedCacheManager::PagedCacheManager(std::shared_ptr model) : CacheManager(model), params_(std::make_shared(*model_)), @@ -187,7 +194,7 @@ void PagedCacheManager::Allocate(const std::vector>& re void PagedCacheManager::Step() { for (auto& request : cache_allocated_requests_) { - if (request->status_ == RequestStatus::Completed) { + if (IsTurnComplete(request->status_)) { continue; } @@ -237,6 +244,11 @@ std::vector> PagedCacheManager::AllocatedRequests() con return cache_allocated_requests_; } +bool PagedCacheManager::IsResident(const std::shared_ptr& request) const { + return std::find(cache_allocated_requests_.begin(), cache_allocated_requests_.end(), request) != + cache_allocated_requests_.end(); +} + std::unique_ptr PagedCacheManager::ReserveStep(const StepPlan& plan) { return std::make_unique( *key_value_cache_, cache_allocated_requests_, plan); diff --git a/src/engine/cache_manager.h b/src/engine/cache_manager.h index f6e382e7b9..3dabd9067f 100644 --- a/src/engine/cache_manager.h +++ b/src/engine/cache_manager.h @@ -54,6 +54,10 @@ struct CacheManager { virtual std::vector> AllocatedRequests() const = 0; + virtual bool IsResident(const std::shared_ptr& request) const = 0; + + virtual size_t ResidentRequestCount() const = 0; + // Columns in the block table the model will see this step, or 0 when the cache does not use one. // The decode path multiplies it by the block size to get the KV length bound it reports through // `attention_metadata`. @@ -97,6 +101,10 @@ struct StaticCacheManager : CacheManager { std::vector> AllocatedRequests() const override; + bool IsResident(const std::shared_ptr& request) const override; + + size_t ResidentRequestCount() const override { return cache_allocated_requests_.size(); } + private: std::shared_ptr params_; std::unique_ptr key_value_cache_; @@ -125,6 +133,10 @@ struct PagedCacheManager : CacheManager { std::vector> AllocatedRequests() const override; + bool IsResident(const std::shared_ptr& request) const override; + + size_t ResidentRequestCount() const override { return cache_allocated_requests_.size(); } + size_t BlockTableColumns() const override { return key_value_cache_->BlockTableColumns(); } size_t MaxQueryTokensPerRequest() const override { diff --git a/src/engine/decoders/simple_decoder.cpp b/src/engine/decoders/simple_decoder.cpp index cc6de1e56c..c816d19fe4 100644 --- a/src/engine/decoders/simple_decoder.cpp +++ b/src/engine/decoders/simple_decoder.cpp @@ -26,7 +26,7 @@ bool IsPureDecodeStep(ScheduledRequests& scheduled_requests) { return false; } for (auto& request : scheduled_requests) { - if (request->IsPrefill() || request->UnprocessedTokens().size() != 1) { + if (request->IsPrefill() || request->ScheduledTokenCount() != 1) { return false; } } diff --git a/src/engine/decoders/static_batch_decoder_io.cpp b/src/engine/decoders/static_batch_decoder_io.cpp index 9843a751f5..4279660ece 100644 --- a/src/engine/decoders/static_batch_decoder_io.cpp +++ b/src/engine/decoders/static_batch_decoder_io.cpp @@ -32,10 +32,10 @@ void StaticBatchDecoderIO::PrepareInputIds(std::shared_ptr mo std::max_element( scheduled_requests.begin(), scheduled_requests.end(), [](const std::shared_ptr& a, const std::shared_ptr& b) { - return a->UnprocessedTokens().size() < b->UnprocessedTokens().size(); + return a->ScheduledTokenCount() < b->ScheduledTokenCount(); }); - const size_t max_sequence_length = (*request_with_max_sequence_length)->UnprocessedTokens().size(); + const size_t max_sequence_length = (*request_with_max_sequence_length)->ScheduledTokenCount(); const size_t batch_size = scheduled_requests.size(); const std::vector input_ids_shape = {static_cast(batch_size), static_cast(max_sequence_length)}; auto input_ids_tensor = std::make_unique(model->p_device_inputs_, Ort::TypeToTensorType); @@ -99,10 +99,10 @@ void StaticBatchDecoderIO::PreparePositionIds(std::shared_ptr std::max_element( scheduled_requests.begin(), scheduled_requests.end(), [](const std::shared_ptr& a, const std::shared_ptr& b) { - return a->UnprocessedTokens().size() < b->UnprocessedTokens().size(); + return a->ScheduledTokenCount() < b->ScheduledTokenCount(); }); - const size_t max_sequence_length = (*request_with_max_sequence_length)->UnprocessedTokens().size(); + const size_t max_sequence_length = (*request_with_max_sequence_length)->ScheduledTokenCount(); const size_t batch_size = scheduled_requests.size(); const std::vector position_ids_shape = {static_cast(batch_size), static_cast(max_sequence_length)}; auto position_ids_tensor = std::make_unique(model->p_device_inputs_, Ort::TypeToTensorType); @@ -113,6 +113,8 @@ void StaticBatchDecoderIO::PreparePositionIds(std::shared_ptr for (size_t i = 0; i < batch_size; ++i) { auto request = scheduled_requests[i]; auto input_ids = request->UnprocessedTokensCpu(); + // A continued request resumes at the sequence length it has already committed to the cache, so + // positions must be offset by it. Only a first prefill starts at zero. const int64_t base_position = request->ProcessedSequenceLength(); for (size_t j = 0; j < max_sequence_length; ++j) { @@ -132,10 +134,10 @@ void StaticBatchDecoderIO::PrepareLogits(std::shared_ptr mode std::max_element( scheduled_requests.begin(), scheduled_requests.end(), [](const std::shared_ptr& a, const std::shared_ptr& b) { - return a->UnprocessedTokens().size() < b->UnprocessedTokens().size(); + return a->ScheduledTokenCount() < b->ScheduledTokenCount(); }); - const int64_t max_sequence_length = (*request_with_max_sequence_length)->UnprocessedTokens().size(); + const int64_t max_sequence_length = (*request_with_max_sequence_length)->ScheduledTokenCount(); const int64_t batch_size = scheduled_requests.size(); const std::vector logits_shape = {batch_size, max_sequence_length, model->config_->model.vocab_size}; logits_ = std::make_unique(model->p_device_inputs_, model->session_info_.GetOutputDataType(model->config_->model.decoder.outputs.logits)); @@ -148,7 +150,12 @@ void StaticBatchDecoderIO::PrepareLogits(std::shared_ptr mode std::vector> StaticBatchDecoderIO::ProcessLogits() { std::vector valid_token_indices; for (auto& request : scheduled_requests_) { - valid_token_indices.push_back(request->UnprocessedTokens().size() - 1); + // A completed row retained in the static batch for continuation can contribute zero + // unprocessed tokens. Selecting index 0 keeps the subspan in bounds; its logits are discarded + // because ScheduledRequests skips completed and removed rows during sampling. + const auto unprocessed_token_count = request->ScheduledTokenCount(); + valid_token_indices.push_back( + unprocessed_token_count == 0 ? 0 : static_cast(unprocessed_token_count - 1)); } // [batch_size, max_sequence_length, vocab_size] diff --git a/src/engine/decoders/varlen_decoder_io.cpp b/src/engine/decoders/varlen_decoder_io.cpp index ff3d2c4a1e..42461d80a9 100644 --- a/src/engine/decoders/varlen_decoder_io.cpp +++ b/src/engine/decoders/varlen_decoder_io.cpp @@ -73,7 +73,7 @@ void VarlenDecoderIO::PrepareInputIds(std::shared_ptr model, plan ? plan->token_count : std::accumulate(scheduled_requests.begin(), scheduled_requests.end(), size_t{0}, [](size_t sum, const std::shared_ptr& request) { - return sum + request->UnprocessedTokens().size(); + return sum + request->ScheduledTokenCount(); }); // On a capturable step the tensors are views onto buffers that were allocated once, so their // device addresses match the ones recorded in the graph. Otherwise each step allocates its own. @@ -189,7 +189,7 @@ void VarlenDecoderIO::PrepareAttentionMetadata(std::shared_ptr(request->UnprocessedTokens().size()); + const int32_t query_len = static_cast(request->ScheduledTokenCount()); // KV length after the step is past length plus query length, which is the current length. const int32_t kv_len = static_cast(request->CurrentSequenceLength()); max_query_len = std::max(max_query_len, query_len); @@ -214,7 +214,7 @@ void VarlenDecoderIO::PrepareLogits(std::shared_ptr model, Sc plan ? plan->token_count : std::accumulate(scheduled_requests.begin(), scheduled_requests.end(), size_t{0}, [](size_t sum, const std::shared_ptr& request) { - return sum + request->UnprocessedTokens().size(); + return sum + request->ScheduledTokenCount(); }); const std::vector logits_shape = {static_cast(num_tokens), static_cast(model->config_->model.vocab_size)}; if (graph_buffers_ != nullptr) { @@ -245,8 +245,8 @@ std::vector> VarlenDecoderIO::ProcessLogits() { } } else { for (size_t i = 0, running_length = 0; i < scheduled_requests_.size(); ++i) { - valid_token_indices[i] = running_length + scheduled_requests_[i]->UnprocessedTokens().size() - 1; - running_length += scheduled_requests_[i]->UnprocessedTokens().size(); + valid_token_indices[i] = running_length + scheduled_requests_[i]->ScheduledTokenCount() - 1; + running_length += scheduled_requests_[i]->ScheduledTokenCount(); } } diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp index 043a75f0a3..367f54060a 100644 --- a/src/engine/engine.cpp +++ b/src/engine/engine.cpp @@ -58,12 +58,51 @@ void Engine::AddRequest(std::shared_ptr request) { if (cache_manager_->SupportsDynamicBatching()) { request->ValidateEngineCompatibility(); } + scheduler_->ValidateRequest(*request); request->Assign(shared_from_this()); scheduler_->AddRequest(request); } void Engine::RemoveRequest(std::shared_ptr request) { + if (request && IsClosed(request->status_)) { + throw std::runtime_error("Cannot remove a request that is already closed."); + } + if (!request || request->engine_.lock().get() != this) { + throw std::runtime_error("Cannot remove a request from an engine it does not belong to."); + } + scheduler_->RemoveRequest(request); + + auto first_undrained = + ready_requests_.begin() + static_cast(ready_request_index_); + const auto retained_end = + std::remove(first_undrained, ready_requests_.end(), request); + const auto new_end = ready_request_index_ == 0 + ? retained_end + : std::move(first_undrained, retained_end, + ready_requests_.begin()); + ready_requests_.erase(new_end, ready_requests_.end()); + ready_request_index_ = 0; + request->CompleteClose(); +} + +void Engine::ValidateRequestCanContinue(const std::shared_ptr& request) const { + if (health_ == EngineHealth::Unhealthy) { + std::rethrow_exception(fatal_error_); + } + if (request->engine_.lock().get() != this) { + throw std::runtime_error("Cannot continue a request that does not belong to this engine."); + } + + if (!cache_manager_->IsResident(request)) { + throw std::runtime_error("Cannot continue a request whose model state is no longer resident."); + } + + if (!cache_manager_->SupportsDynamicBatching() && + cache_manager_->ResidentRequestCount() > 1) { + throw std::runtime_error( + "Continuous decoding is only supported when a static engine batch contains one request."); + } } std::shared_ptr Engine::Step() { @@ -79,11 +118,22 @@ std::shared_ptr Engine::Step() { std::shared_ptr Engine::StepStatic() { while (scheduler_->HasPendingRequests()) { auto scheduled_requests = scheduler_->Schedule(); + std::vector statuses_before_step; + statuses_before_step.reserve(scheduled_requests.size()); + for (const auto& request : scheduled_requests) { + statuses_before_step.push_back(request->status_); + } + model_executor_->Decode(scheduled_requests); scheduled_requests.GenerateNextTokens(); - for (auto& request : scheduled_requests) { - if (request->HasUnseenTokens() || request->IsDone()) { + for (size_t i = 0; i < scheduled_requests.size(); ++i) { + auto request = scheduled_requests[i]; + const bool turn_completed_this_step = + !IsTurnComplete(statuses_before_step[i]) && + IsTurnComplete(request->status_); + if (!IsClosed(request->status_) && + (request->HasUnseenTokens() || turn_completed_this_step)) { ready_requests_.push_back(request); } } @@ -146,8 +196,31 @@ std::shared_ptr Engine::StepDynamic() { std::current_exception()); } - auto scheduled_requests = - scheduler_->CreateScheduledRequests(step_plan_); + auto scheduled_requests = [&]() -> ScheduledRequests { + try { + return scheduler_->CreateScheduledRequests(step_plan_); + } catch (...) { + const auto construction_error = std::current_exception(); + try { + reservation->Release(); + } catch (...) { + ++transaction_metrics_.rollbacks; + MarkUnhealthyAndThrow( + StepOutcomeKind::FatalExecutionFailure, + step_plan_.transaction_id, + nullptr, + "Failed to release cache state after scheduled-request construction failed.", + std::current_exception()); + } + ++transaction_metrics_.rollbacks; + MarkUnhealthyAndThrow( + StepOutcomeKind::ExecutionContractFailure, + step_plan_.transaction_id, + nullptr, + "Failed to construct the scheduled request transaction.", + construction_error); + } + }(); ExecutionContext context{&step_plan_}; context.cache_reservation = reservation->PagedReservation(); diff --git a/src/engine/engine.h b/src/engine/engine.h index a98def57d1..0e734e7f3d 100644 --- a/src/engine/engine.h +++ b/src/engine/engine.h @@ -122,6 +122,7 @@ struct Engine : std::enable_shared_from_this, std::shared_ptr DrainReadyRequest(); std::shared_ptr StepDynamic(); std::shared_ptr StepStatic(); + void ValidateRequestCanContinue(const std::shared_ptr& request) const; [[noreturn]] void MarkUnhealthyAndThrow(StepOutcomeKind outcome, StepTransactionId transaction_id, const void* request_id, @@ -141,6 +142,8 @@ struct Engine : std::enable_shared_from_this, std::vector> ready_requests_; std::vector> staged_ready_requests_; size_t ready_request_index_{}; + + friend struct Request; }; } // namespace Generators diff --git a/src/engine/engine_invariants.h b/src/engine/engine_invariants.h index 53cbc900cd..164c56ed81 100644 --- a/src/engine/engine_invariants.h +++ b/src/engine/engine_invariants.h @@ -35,7 +35,7 @@ struct RequestStateSnapshot { RequestStatus status{RequestStatus::Unassigned}; int64_t current_sequence_length{}; // Total tokens the Request's search currently holds. int64_t processed_sequence_length{}; // Tokens the model has already processed into the cache. - int64_t seen_sequence_length{}; // Tokens already streamed to (seen by) the application. + int64_t seen_sequence_length{}; // High-water sequence index of consumed generated output. bool is_prefill{}; }; diff --git a/src/engine/request.cpp b/src/engine/request.cpp index 9434c67e59..d9be5d1dc3 100644 --- a/src/engine/request.cpp +++ b/src/engine/request.cpp @@ -20,6 +20,18 @@ DeviceSpan AllocateOnDevice(GeneratorParams& params, return device_tokens; } +void ValidateAppendLength(const GeneratorParams& params, + size_t current_sequence_length, + size_t token_count) { + const size_t max_length = static_cast(params.search.max_length); + if (current_sequence_length >= max_length || + token_count >= max_length - current_sequence_length) { + throw std::runtime_error( + "Input tokens must leave room for at least one generated token before max_length (" + + std::to_string(params.search.max_length) + ")."); + } +} + } // namespace Request::Request(std::shared_ptr params) @@ -57,6 +69,7 @@ void Request::Assign(std::shared_ptr engine) { prompt_sequence_length_ = CurrentSequenceLength(); seen_sequence_length_ = CurrentSequenceLength(); tokens_host_.reserve(params_->search.max_length); + unseen_token_indices_.reserve(params_->search.max_length); tokens_host_.insert(tokens_host_.end(), prefill_input_ids_.begin(), prefill_input_ids_.end()); prefill_input_ids_.clear(); } @@ -74,33 +87,89 @@ void Request::Schedule() { } void Request::Remove() { + if (status_ == RequestStatus::Unassigned) { + throw std::runtime_error("Cannot close a request that has not been submitted to an engine."); + } + if (IsClosed(status_)) { + throw std::runtime_error("Cannot close a request that is already closed."); + } + auto engine = engine_.lock(); - if (engine) { - engine->RemoveRequest(shared_from_this()); + if (!engine) { + CompleteClose(); + return; } - status_ = RequestStatus::Unassigned; + engine->RemoveRequest(shared_from_this()); +} + +void Request::CompleteClose() { + engine_.reset(); + status_ = RequestStatus::Closed; } void Request::AddTokens(std::span tokens) { - if (tokens.size() == 0) + if (tokens.empty()) throw std::runtime_error("Expected at least one token for generation. Received 0."); - if (tokens.size() + CurrentSequenceLength() > params_->search.max_length) - throw std::runtime_error("Input tokens size (" + - std::to_string(tokens.size()) + - ") exceeds the max length (" + - std::to_string(params_->search.max_length) + ")"); + if (status_ != RequestStatus::Unassigned) { + if (IsTurnComplete(status_)) { + throw std::runtime_error("AddTokens only accepts initial input; use Continue for another turn."); + } + if (IsClosed(status_)) { + throw std::runtime_error("Cannot add tokens to a closed request."); + } + throw std::runtime_error("AddTokens only accepts initial input before submission to an engine."); + } + + ValidateAppendLength(*params_, prefill_input_ids_.size(), tokens.size()); + std::copy(tokens.begin(), tokens.end(), std::back_inserter(prefill_input_ids_)); +} - if (status_ == RequestStatus::Unassigned) { - std::copy(tokens.begin(), tokens.end(), std::back_inserter(prefill_input_ids_)); - } else if (status_ == RequestStatus::InProgress) { - throw std::runtime_error("Cannot add tokens to a request that is in progress."); - } else if (status_ == RequestStatus::Completed) { - auto device_tokens = AllocateOnDevice(*params_, tokens); +void Request::Continue(std::span tokens) { + if (tokens.empty()) + throw std::runtime_error("Expected at least one token for continuation. Received 0."); + if (!IsTurnComplete(status_)) { + if (IsClosed(status_)) { + throw std::runtime_error("Cannot continue a closed request."); + } + throw std::runtime_error("Continue is only valid after the current turn is complete."); + } + + auto engine = engine_.lock(); + if (!engine) { + throw std::runtime_error("Cannot continue a request after its engine has been destroyed."); + } + const DeviceType cache_device = params_->model_->p_device_kvcache_->GetType(); + if (!SupportsContinuousDecoding(cache_device)) { + throw std::runtime_error( + "Continuous decoding is not supported on the selected KV-cache device type (" + + to_string(cache_device) + ")."); + } + engine->ValidateRequestCanContinue(shared_from_this()); + ValidateAppendLength(*params_, static_cast(CurrentSequenceLength()), tokens.size()); + if (tokens_host_.capacity() < tokens_host_.size() + tokens.size()) { + throw std::logic_error("The request host token mirror does not have reserved continuation capacity."); + } + + auto device_tokens = AllocateOnDevice(*params_, tokens); + search_->SaveStateForTransaction(); + try { search_->AppendTokens(device_tokens); - prompt_sequence_length_ = CurrentSequenceLength(); - tokens_host_.insert(tokens_host_.end(), tokens.begin(), tokens.end()); + search_->CommitStateForTransaction(); + } catch (...) { + const auto append_error = std::current_exception(); + try { + search_->RestoreStateForTransaction(); + } catch (...) { + throw std::runtime_error( + "Continue failed and the request search state could not be restored."); + } + std::rethrow_exception(append_error); } + + tokens_host_.insert(tokens_host_.end(), tokens.begin(), tokens.end()); + prompt_sequence_length_ = CurrentSequenceLength(); + status_ = RequestStatus::Assigned; } int64_t Request::CurrentSequenceLength() const { @@ -152,14 +221,23 @@ void Request::AdvanceChunk() { } int32_t Request::UnseenToken() { - if (static_cast(seen_sequence_length_) >= tokens_host_.size()) + if (next_unseen_token_index_ == unseen_token_indices_.size()) throw std::runtime_error("All tokens have been seen."); - return tokens_host_[seen_sequence_length_++]; + const size_t token_index = unseen_token_indices_[next_unseen_token_index_++]; + if (token_index >= tokens_host_.size()) + throw std::runtime_error("The unseen token index is outside the host token sequence."); + seen_sequence_length_ = std::max(seen_sequence_length_, static_cast(token_index + 1)); + const int32_t token = tokens_host_[token_index]; + if (next_unseen_token_index_ == unseen_token_indices_.size()) { + unseen_token_indices_.clear(); + next_unseen_token_index_ = 0; + } + return token; } bool Request::HasUnseenTokens() const { - return seen_sequence_length_ < CurrentSequenceLength(); + return next_unseen_token_index_ < unseen_token_indices_.size(); } DeviceSpan Request::UnprocessedTokens() { @@ -177,7 +255,7 @@ std::span Request::UnprocessedTokensCpu() const { } bool Request::IsDone() const { - return status_ == RequestStatus::Completed; + return status_ == RequestStatus::TurnComplete; } bool Request::IsPrefill() const { @@ -268,10 +346,12 @@ void Request::CommitStateForTransaction() { void Request::CommitStep(const RequestStepPlan& plan, const RequestStepResult& result) noexcept { if (result.token_appended) { + const size_t token_index = tokens_host_.size(); tokens_host_.push_back(result.token); + unseen_token_indices_.push_back(token_index); } processed_sequence_length_ = static_cast(plan.target_cache_slots); - status_ = result.done ? RequestStatus::Completed : RequestStatus::InProgress; + status_ = result.done ? RequestStatus::TurnComplete : RequestStatus::InProgress; } void Request::ApplyLogitsProcessors(DeviceSpan logits) { @@ -349,11 +429,15 @@ void Request::CompleteGeneration() { if (new_token_count > next_tokens.size()) throw std::runtime_error("The search produced fewer tokens than it appended to the sequence."); + const size_t first_new_token = tokens_host_.size(); tokens_host_.insert(tokens_host_.end(), next_tokens.end() - new_token_count, next_tokens.end()); + for (size_t token_index = first_new_token; token_index < tokens_host_.size(); ++token_index) { + unseen_token_indices_.push_back(token_index); + } } if (search_->IsDone()) { - status_ = RequestStatus::Completed; + status_ = RequestStatus::TurnComplete; } } diff --git a/src/engine/request.h b/src/engine/request.h index 1a009232f2..fdf933f54a 100644 --- a/src/engine/request.h +++ b/src/engine/request.h @@ -55,11 +55,24 @@ struct Request : std::enable_shared_from_this, void Schedule(); /** - * @brief Adds a sequence of tokens to the request for processing. + * @brief Adds initial input tokens before the request is submitted to an Engine. * @param tokens Span of token IDs to be added. + * + * This operation is legal only while the request is Unassigned. Use Continue() + * to begin another turn after the current turn reaches TurnComplete. */ void AddTokens(std::span tokens); + /** + * @brief Queues another generation turn using resident model state. + * @param tokens New input tokens to append after the completed turn. + * + * This operation is legal only from TurnComplete. It preserves unread generated + * output, appends no input tokens to that output stream, and moves the request + * back to Assigned (the queued state). + */ + void Continue(std::span tokens); + /** * @brief Retrieves the next unseen token in the request. * @return The next unseen token ID. @@ -87,7 +100,7 @@ struct Request : std::enable_shared_from_this, /** * @brief Returns the unprocessed tokens from the host-side mirror of the sequence. * @return Span of unprocessed token IDs, valid only until the next call that appends to the - * sequence (CompleteGeneration, AddTokens or Assign). Copy it if it must outlive those. + * sequence (CompleteGeneration, Continue or Assign). Copy it if it must outlive those. * * Same tokens as UnprocessedTokens(), but readable without copying them back from the device. * Building the next step's input ids is the hot path for this, and a device readback there costs @@ -133,13 +146,15 @@ struct Request : std::enable_shared_from_this, void CompleteGeneration(); /** - * @brief Checks if the termination condition for the request has been met. - * @return True if the request is done, false otherwise. + * @brief Checks if the current generation turn reached a stopping condition. + * @return True in TurnComplete; the request may still be continued or closed. */ bool IsDone() const; + RequestStatus Status() const noexcept { return status_; } + /** - * @brief Removes the request from being processed. + * @brief Removes the request from its engine and moves it to terminal Closed. */ void Remove(); @@ -189,6 +204,13 @@ struct Request : std::enable_shared_from_this, */ void BindScheduledTokenCount(size_t token_count); + /** + * @brief Tokens this request contributes to the next step. + * + * Equivalent to UnprocessedTokens().size() without constructing a device span. + */ + size_t ScheduledTokenCount() const; + /** * @brief True when this step's tokens run to the end of the sequence. * @@ -263,17 +285,21 @@ struct Request : std::enable_shared_from_this, void* GetOpaqueData(); private: - // Tokens of the current step, clamped to what is actually left to process. - size_t ScheduledTokenCount() const; - // The search sequence is partitioned at processed_sequence_length_: tokens before it already // have KV entries, and UnprocessedTokens() returns the scheduled prefix of [processed, current). - // seen_sequence_length_ independently tracks tokens consumed by the application. + // seen_sequence_length_ is the high-water sequence index of generated output consumed by the + // application. Continuation input creates gaps, so it is not an unseen-token count. std::vector prefill_input_ids_; // Host-side mirror of the full sequence (prompt + generated tokens). Kept in step with the // search's device sequence so that streaming and input-id preparation never read it back. std::vector tokens_host_; + std::vector unseen_token_indices_; + size_t next_unseen_token_index_{}; int64_t seen_sequence_length_{}; + friend struct Engine; + + void CompleteClose(); + int64_t processed_sequence_length_{}; // Sequence length the application's tokens reach up to. Everything below it is prompt, so the // request is still prefilling while processed_sequence_length_ has not caught up with it. diff --git a/src/engine/request_status.h b/src/engine/request_status.h index 5b6859ef27..8809c30173 100644 --- a/src/engine/request_status.h +++ b/src/engine/request_status.h @@ -15,11 +15,31 @@ namespace Generators { enum class RequestStatus { - Unassigned, // A request has been created but has not been added to the engine yet. - // This is the state of a request when it is first created. - Assigned, // The request has been added to the engine and is waiting to be scheduled. - InProgress, // The request has been scheduled and is currently being processed. - Completed, // The request has been completed successfully. + Unassigned, // Created: initial input may be added before submission to an Engine. + Assigned, // Queued: submitted initial work or a resident continuation awaits execution. + InProgress, // The current generation turn is executable and owned by the Engine. + TurnComplete, // The current turn stopped; output and resident model state remain available. + Closed, // Permanently terminal; no scheduler or cache resources remain owned. }; +constexpr bool IsQueued(RequestStatus status) noexcept { + return status == RequestStatus::Assigned; +} + +constexpr bool IsExecuting(RequestStatus status) noexcept { + return status == RequestStatus::InProgress; +} + +constexpr bool IsExecutable(RequestStatus status) noexcept { + return IsQueued(status) || IsExecuting(status); +} + +constexpr bool IsTurnComplete(RequestStatus status) noexcept { + return status == RequestStatus::TurnComplete; +} + +constexpr bool IsClosed(RequestStatus status) noexcept { + return status == RequestStatus::Closed; +} + } // namespace Generators diff --git a/src/engine/scheduled_requests.cpp b/src/engine/scheduled_requests.cpp index 6dd7843ad0..c3cb824b0e 100644 --- a/src/engine/scheduled_requests.cpp +++ b/src/engine/scheduled_requests.cpp @@ -56,6 +56,9 @@ ScheduledRequests::ScheduledRequests(const StepPlan& plan, request_ids.end()) { throw std::runtime_error("The dynamic step plan contains an invalid request."); } + if (!IsExecutable(entry.request->status_)) { + throw std::runtime_error("The dynamic step plan contains a request that is not executable."); + } const int64_t remaining = entry.request->CurrentSequenceLength() - entry.request->ProcessedSequenceLength(); @@ -105,21 +108,22 @@ void ScheduledRequests::GenerateNextTokens() { // serialize the whole batch; launching all of them first means only the first completion below // actually waits for the device. for (size_t request_idx = 0; request_idx < requests_.size(); ++request_idx) { - if (requests_[request_idx]->status_ != RequestStatus::Completed && + if (IsExecuting(requests_[request_idx]->status_) && requests_[request_idx]->IsChunkComplete()) { requests_[request_idx]->GenerateNextTokens(logits[request_idx]); } } for (size_t request_idx = 0; request_idx < requests_.size(); ++request_idx) { - if (requests_[request_idx]->status_ != RequestStatus::Completed && + if (IsExecuting(requests_[request_idx]->status_) && requests_[request_idx]->IsChunkComplete()) { requests_[request_idx]->CompleteGeneration(); } } for (const auto& request : requests_) { - if (request->status_ != RequestStatus::Completed && !request->IsChunkComplete()) + if (IsExecuting(request->status_) && + !request->IsChunkComplete()) request->AdvanceChunk(); } } catch (...) { @@ -152,7 +156,7 @@ bool ScheduledRequests::TryGenerateNextTokensBatched(std::vectorstatus_ != RequestStatus::Completed && + if (IsExecuting(requests_[request_idx]->status_) && requests_[request_idx]->IsChunkComplete()) sampling_plan_->logits.push_back(logits[request_idx]); } @@ -180,7 +184,8 @@ bool ScheduledRequests::TryGenerateNextTokensBatched(std::vectorCompleteGeneration(); } for (const auto& request : requests_) { - if (request->status_ != RequestStatus::Completed && !request->IsChunkComplete()) + if (IsExecuting(request->status_) && + !request->IsChunkComplete()) request->AdvanceChunk(); } @@ -196,7 +201,13 @@ bool ScheduledRequests::PrepareBatchedSamplingPlan( sampling_plan_->Clear(); for (const auto& request : requests_) { - if (request->status_ == RequestStatus::Completed || !request->IsChunkComplete()) + // Dynamic transactions keep newly admitted and continued requests Queued until commit, while + // the static scheduler moves every executable row to InProgress before constructing the batch. + const bool status_is_executable = + require_transaction_support ? IsExecutable(request->status_) + : IsExecuting(request->status_); + if (!status_is_executable || + !request->IsChunkComplete()) continue; const auto args = ResolveSampleArgs(request->SearchOptions()); diff --git a/src/engine/scheduler.cpp b/src/engine/scheduler.cpp index 04f20c7c2a..2315066b14 100644 --- a/src/engine/scheduler.cpp +++ b/src/engine/scheduler.cpp @@ -31,13 +31,16 @@ ScheduledRequests Scheduler::CreateScheduledRequests(const StepPlan& plan) { StaticBatchScheduler::StaticBatchScheduler(std::shared_ptr model, std::shared_ptr cache_manager) : Scheduler{model}, model_{model}, cache_manager_{cache_manager} {} -void StaticBatchScheduler::AddRequest(std::shared_ptr request) { +void StaticBatchScheduler::ValidateRequest(const Request& request) const { // The static batch decoder rebuilds its contiguous cache from the whole sequence every step, so it // cannot resume a half written prompt. Only the paged cache can hold one. - if (request->SearchOptions().chunk_size.value_or(0) != 0) { + if (request.SearchOptions().chunk_size.value_or(0) != 0) { throw std::runtime_error( "search.chunk_size requires dynamic batching; the static batch scheduler cannot chunk a prefill."); } +} + +void StaticBatchScheduler::AddRequest(std::shared_ptr request) { if (auto* sampler = GetBatchedSampler()) request->SamplingState(*sampler); requests_pool_.push_back(request); @@ -47,15 +50,29 @@ void StaticBatchScheduler::RemoveRequest(std::shared_ptr request) { // For statically batched requests, memory is managed as a single block for the entire batch, // so individual requests cannot be deallocated until the whole batch is completed. // Therefore, deallocation is only performed for dynamically batched requests below. - // we simply mark the request to be removed and it will be deallocated when the - // entire batch is completed. - to_be_removed_requests_.insert(request); + if (!cache_manager_->IsResident(request)) { + requests_pool_.erase( + std::remove(requests_pool_.begin(), requests_pool_.end(), request), + requests_pool_.end()); + } } ScheduledRequests StaticBatchScheduler::Schedule() { + const auto allocated_requests = cache_manager_->AllocatedRequests(); + const auto is_resident = [&allocated_requests](const std::shared_ptr& request) { + return std::find(allocated_requests.begin(), allocated_requests.end(), request) != + allocated_requests.end(); + }; + + for (const auto& request : allocated_requests) { + if (IsQueued(request->status_)) { + request->Schedule(); + } + } + std::vector> requests_to_schedule; for (auto& request : requests_pool_) { - if (request->status_ == RequestStatus::Assigned) { + if (IsQueued(request->status_) && !is_resident(request)) { requests_to_schedule.push_back(request); } } @@ -67,12 +84,8 @@ ScheduledRequests StaticBatchScheduler::Schedule() { requests_to_schedule.begin() + batch_size); if (cache_manager_->CanAllocate(batch_requests)) { // Before allocating, we need to ensure that the existing requests in the cache manager - // are complete and that if they were previously removed from the engine, they are no longer - // in the requests pool. - for (auto& request : cache_manager_->AllocatedRequests()) { - if (request->status_ != RequestStatus::Completed && to_be_removed_requests_.count(request)) { - throw std::runtime_error("Encountered a request that was removed from the engine but was not completed."); - } + // are terminal and no longer need to remain in the scheduler pool. + for (auto& request : allocated_requests) { requests_pool_.erase(std::remove(requests_pool_.begin(), requests_pool_.end(), request), requests_pool_.end()); } @@ -97,7 +110,7 @@ ScheduledRequests StaticBatchScheduler::Schedule() { bool StaticBatchScheduler::HasPendingRequests() const { for (auto& request : requests_pool_) { - if (request->status_ != RequestStatus::Completed) { + if (IsExecutable(request->status_)) { return true; } } @@ -125,30 +138,7 @@ ScheduledRequests DynamicBatchScheduler::Schedule() { "Dynamic batching requires transactional step planning."); } -void DynamicBatchScheduler::ReapCompletedRequests() { - auto allocated_requests = cache_manager_->AllocatedRequests(); - std::vector> completed_requests; - std::copy_if(allocated_requests.begin(), allocated_requests.end(), - std::back_inserter(completed_requests), - [](const std::shared_ptr& request) { - return request->status_ == RequestStatus::Completed; - }); - if (!completed_requests.empty()) { - cache_manager_->Deallocate(completed_requests); - requests_pool_.erase( - std::remove_if(requests_pool_.begin(), requests_pool_.end(), - [](const std::shared_ptr& request) { - return request->status_ == RequestStatus::Completed; - }), - requests_pool_.end()); - } -} - StepPlanningResult DynamicBatchScheduler::PlanStep(StepPlan& plan) { - // Completed requests release their blocks before admission, making that capacity available to - // requests waiting in Assigned state during this same planning pass. - ReapCompletedRequests(); - plan.requests.clear(); plan.scheduled_request_limit = 0; plan.token_count = 0; @@ -160,16 +150,18 @@ StepPlanningResult DynamicBatchScheduler::PlanStep(StepPlan& plan) { DecodeFirstBudgetCandidate budget; size_t processed_sequence_length{}; }; + const auto allocated_requests = cache_manager_->AllocatedRequests(); std::vector candidates; + candidates.reserve(allocated_requests.size() + requests_pool_.size()); const size_t cache_query_token_cap = cache_manager_->MaxQueryTokensPerRequest(); const auto add_candidate = [&candidates, cache_query_token_cap]( const std::shared_ptr& request, bool newly_admitted) { const auto snapshot = request->Snapshot(); - const RequestStatus expected_status = - newly_admitted ? RequestStatus::Assigned : RequestStatus::InProgress; - if (snapshot.status != expected_status) { + const bool valid_status = + newly_admitted ? IsQueued(snapshot.status) : IsExecutable(snapshot.status); + if (!valid_status) { throw std::runtime_error("Request status is invalid for dynamic step planning."); } const auto remaining_token_count = @@ -204,13 +196,19 @@ StepPlanningResult DynamicBatchScheduler::PlanStep(StepPlan& plan) { candidates.push_back(std::move(candidate)); }; - const auto allocated_requests = cache_manager_->AllocatedRequests(); + const auto is_resident = [&allocated_requests](const std::shared_ptr& request) { + return std::find(allocated_requests.begin(), allocated_requests.end(), request) != + allocated_requests.end(); + }; for (const auto& request : allocated_requests) { + if (IsTurnComplete(request->status_)) { + continue; + } add_candidate(request, false); } for (const auto& request : requests_pool_) { - if (request->status_ == RequestStatus::Assigned) { + if (IsQueued(request->status_) && !is_resident(request)) { add_candidate(request, true); } } @@ -236,7 +234,9 @@ StepPlanningResult DynamicBatchScheduler::PlanStep(StepPlan& plan) { } std::vector selected_candidates; + std::vector selected_processed_lengths; selected_candidates.reserve(plan.requests.size()); + selected_processed_lengths.reserve(plan.requests.size()); for (const auto& entry : plan.requests) { const auto candidate = std::find_if( candidates.begin(), candidates.end(), @@ -246,6 +246,7 @@ StepPlanningResult DynamicBatchScheduler::PlanStep(StepPlan& plan) { if (candidate == candidates.end()) throw std::logic_error("Cache planning selected an unknown request."); selected_candidates.push_back(candidate->budget); + selected_processed_lengths.push_back(candidate->processed_sequence_length); } const auto token_counts = AllocateDecodeFirstTokenBudget( selected_candidates, dynamic_batching.max_scheduled_tokens); @@ -257,16 +258,9 @@ StepPlanningResult DynamicBatchScheduler::PlanStep(StepPlan& plan) { plan.graph_capture_eligible = true; for (size_t i = 0; i < plan.requests.size(); ++i) { auto& entry = plan.requests[i]; - const auto candidate = std::find_if( - candidates.begin(), candidates.end(), - [&entry](const Candidate& value) { - return value.entry.request_id == entry.request_id; - }); - if (candidate == candidates.end()) - throw std::logic_error("Cache planning selected an unknown request."); entry.unprocessed_token_count = token_counts[i]; entry.target_cache_slots = RequiredSlots( - candidate->processed_sequence_length, + selected_processed_lengths[i], entry.unprocessed_token_count); entry.packed_token_offset = packed_token_offset; entry.logits_row_index = @@ -281,7 +275,7 @@ StepPlanningResult DynamicBatchScheduler::PlanStep(StepPlan& plan) { bool DynamicBatchScheduler::HasPendingRequests() const { for (auto& request : requests_pool_) { - if (request->status_ != RequestStatus::Completed) { + if (IsExecutable(request->status_)) { return true; } } diff --git a/src/engine/scheduler.h b/src/engine/scheduler.h index b008d611e7..5e1080d03b 100644 --- a/src/engine/scheduler.h +++ b/src/engine/scheduler.h @@ -26,6 +26,8 @@ struct Scheduler { static std::unique_ptr Create(std::shared_ptr model, std::shared_ptr cache_manager); + virtual void ValidateRequest(const Request&) const {} + /** * @brief Adds a request to the Scheduler for processing. * @param request A shared pointer to the Request object to be added. @@ -82,6 +84,8 @@ struct Scheduler { struct StaticBatchScheduler : Scheduler { StaticBatchScheduler(std::shared_ptr model, std::shared_ptr cache_manager); + void ValidateRequest(const Request& request) const override; + void AddRequest(std::shared_ptr request) override; void RemoveRequest(std::shared_ptr request) override; @@ -94,7 +98,6 @@ struct StaticBatchScheduler : Scheduler { std::shared_ptr model_; std::shared_ptr cache_manager_; std::vector> requests_pool_; - std::set> to_be_removed_requests_; }; struct DynamicBatchScheduler : Scheduler { @@ -111,8 +114,6 @@ struct DynamicBatchScheduler : Scheduler { bool HasPendingRequests() const override; private: - void ReapCompletedRequests(); - std::shared_ptr model_; std::shared_ptr cache_manager_; std::vector> requests_pool_; diff --git a/src/generators.cpp b/src/generators.cpp index 07d65c895e..baaa985377 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -59,6 +59,20 @@ namespace Generators { static bool _ = (Ort::InitApi(), false); +bool SupportsContinuousDecoding(DeviceType device_type) noexcept { + // Some models fall back to CPU attention, where continuation is valid because their KV cache is + // also on CPU. Other listed providers preserve appendable KV state across generation turns. + constexpr std::array supported_devices{ + DeviceType::CPU, + DeviceType::CUDA, + DeviceType::WEBGPU, + DeviceType::OpenVINO, + DeviceType::NvTensorRtRtx, + DeviceType::RyzenAI}; + return std::find(supported_devices.begin(), supported_devices.end(), device_type) != + supported_devices.end(); +} + static OrtLoggingLevel GetDefaultOrtLoggingLevel() { bool ort_verbose_logging = false; GetEnv("ORTGENAI_ORT_VERBOSE_LOGGING", ort_verbose_logging); @@ -669,19 +683,8 @@ void Generator::AppendTokens(cpu_span input_ids) { if (search_->GetSequenceLength() != 0 && state_->params_->search.batch_size > 1) throw std::runtime_error("AppendTokens can only be called once for batch_size > 1. To call AppendTokens again, use RewindToLength(0)"); - // Some models fallback to CPU for the attention operator (for example, some decoder-pipeline NPU models). - // Continuous decoding is supported for this case as the kv cache for such models is always on CPU. - constexpr std::array devices_supporting_continuous_decoding{ - DeviceType::CPU, - DeviceType::CUDA, - DeviceType::WEBGPU, - DeviceType::OpenVINO, - DeviceType::NvTensorRtRtx, - DeviceType::RyzenAI}; - if (search_->GetSequenceLength() != 0 && - std::none_of(devices_supporting_continuous_decoding.begin(), devices_supporting_continuous_decoding.end(), - [this](DeviceType device_type) { return device_type == state_->model_.p_device_kvcache_->GetType(); })) + !SupportsContinuousDecoding(state_->model_.p_device_kvcache_->GetType())) // Support for continuous decoding should be based on the type of device used for KV cache throw std::runtime_error("Continuous decoding is not supported on the selected device type (" + to_string(state_->model_.p_device_kvcache_->GetType()) + "). Please recreate the generator instance to avoid using continuous decoding."); diff --git a/src/generators.h b/src/generators.h index 19349bc596..70ab464711 100644 --- a/src/generators.h +++ b/src/generators.h @@ -73,6 +73,7 @@ using TokenSequences = std::vector>; std::string to_string(DeviceType device_type); DeviceInterface* GetDeviceInterface(DeviceType type); +bool SupportsContinuousDecoding(DeviceType device_type) noexcept; struct GeneratorParams : std::enable_shared_from_this, LeakChecked, ExternalRefCounted { GeneratorParams(const Config& config); // This constructor is only used for internal generator benchmarks diff --git a/src/ort_genai.h b/src/ort_genai.h index e3c45b04d5..1e761b8251 100644 --- a/src/ort_genai.h +++ b/src/ort_genai.h @@ -899,12 +899,22 @@ struct OgaRequest : OgaAbstract { OgaCheckResult(OgaRequestAddTokens(this, &tokens)); } + void Continue(const OgaSequences& tokens) { + OgaCheckResult(OgaRequestContinue(this, &tokens)); + } + bool IsDone() const { bool is_done{}; OgaCheckResult(OgaRequestIsDone(this, &is_done)); return is_done; } + OgaRequestStatus GetStatus() const { + OgaRequestStatus status; + OgaCheckResult(OgaRequestGetStatus(this, &status)); + return status; + } + bool HasUnseenTokens() const { bool has_unseen_tokens{}; OgaCheckResult(OgaRequestHasUnseenTokens(this, &has_unseen_tokens)); diff --git a/src/ort_genai_c.cpp b/src/ort_genai_c.cpp index 41ed555f6b..18f5a7b95b 100644 --- a/src/ort_genai_c.cpp +++ b/src/ort_genai_c.cpp @@ -1329,14 +1329,24 @@ OgaResult* OgaCreateRequest(OgaGeneratorParams* params, OgaRequest** out) { OgaResult* OgaRequestAddTokens(OgaRequest* request, const OgaSequences* tokens) { OGA_TRY - if (tokens->size() > 1) { - throw std::runtime_error("Request can only be created with a single sequence"); + if (tokens->size() != 1) { + throw std::runtime_error("Request input must contain exactly one sequence."); } request->AddTokens((*tokens)[0]); return nullptr; OGA_CATCH } +OgaResult* OgaRequestContinue(OgaRequest* request, const OgaSequences* tokens) { + OGA_TRY + if (tokens->size() != 1) { + throw std::runtime_error("Request continuation must contain exactly one sequence."); + } + request->Continue((*tokens)[0]); + return nullptr; + OGA_CATCH +} + OgaResult* OgaRequestHasUnseenTokens(const OgaRequest* request, bool* out) { OGA_TRY *out = request->HasUnseenTokens(); @@ -1358,6 +1368,29 @@ OgaResult* OgaRequestIsDone(const OgaRequest* request, bool* out) { OGA_CATCH } +OgaResult* OgaRequestGetStatus(const OgaRequest* request, OgaRequestStatus* out) { + OGA_TRY + switch (request->Status()) { + case Generators::RequestStatus::Unassigned: + *out = OgaRequestStatus_created; + break; + case Generators::RequestStatus::Assigned: + *out = OgaRequestStatus_queued; + break; + case Generators::RequestStatus::InProgress: + *out = OgaRequestStatus_in_progress; + break; + case Generators::RequestStatus::TurnComplete: + *out = OgaRequestStatus_turn_complete; + break; + case Generators::RequestStatus::Closed: + *out = OgaRequestStatus_closed; + break; + } + return nullptr; + OGA_CATCH +} + OgaResult* OgaRequestSetOpaqueData(OgaRequest* request, void* data) { OGA_TRY request->SetOpaqueData(data); diff --git a/src/ort_genai_c.h b/src/ort_genai_c.h index 264f3e4f51..6bf9ac40f0 100644 --- a/src/ort_genai_c.h +++ b/src/ort_genai_c.h @@ -59,6 +59,14 @@ typedef enum OgaElementType { OgaElementType_bfloat16, // Non-IEEE floating-point format based on IEEE754 single-precision } OgaElementType; +typedef enum OgaRequestStatus { + OgaRequestStatus_created, + OgaRequestStatus_queued, + OgaRequestStatus_in_progress, + OgaRequestStatus_turn_complete, + OgaRequestStatus_closed, +} OgaRequestStatus; + typedef struct OgaResult OgaResult; typedef struct OgaGeneratorParams OgaGeneratorParams; typedef struct OgaGenerator OgaGenerator; @@ -1205,7 +1213,8 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaEngineAddRequest(OgaEngine* engine, OgaReq * \brief Removes a request from the OgaEngine. * * This function removes a request from the engine, allowing it to be cleaned up. The request must have been previously added - * to the engine using OgaEngineAddRequest. After this call, the request will no longer be processed by the engine. + * to the engine using OgaEngineAddRequest. After this call, the request will no longer be processed and cannot be reused. + * Removing an already closed request returns an error. * * \param[in] engine The engine instance from which the request is being removed. * \param[in] request The request to remove from the engine. The request must have been previously added to the engine. @@ -1226,10 +1235,11 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaEngineRemoveRequest(OgaEngine* engine, Oga OGA_EXPORT OgaResult* OGA_API_CALL OgaCreateRequest(OgaGeneratorParams* params, OgaRequest** out); /** - * \brief Adds input sequences to the request. + * \brief Adds initial input sequences to a created request. * - * This function sets the input sequences for the request. The input sequences are used to seed the generation process. - * The request must have been created using OgaCreateRequest before calling this function. + * This function is valid only before the request is submitted to an Engine. Use OgaRequestContinue to begin another + * generation turn after OgaRequestStatus_turn_complete. + * Input must leave room for at least one generated token below max_length. * * \param[in] request The request to set the input sequences on. * \param[in] tokens The input sequences to set on the request. @@ -1237,6 +1247,18 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaCreateRequest(OgaGeneratorParams* params, */ OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestAddTokens(OgaRequest* request, const OgaSequences* tokens); +/** + * \brief Queues another generation turn using the request's resident model state. + * + * This function is valid only from OgaRequestStatus_turn_complete. The request moves to + * OgaRequestStatus_queued, and subsequent OgaEngineStep calls process the appended input. + * + * \param[in] request The request to continue. + * \param[in] tokens The new input sequence for the next turn. + * \return OgaResult containing the error message if continuation failed, or nullptr on success. + */ +OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestContinue(OgaRequest* request, const OgaSequences* tokens); + /** * \brief Destroys the given request. * @@ -1300,11 +1322,10 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestHasUnseenTokens(const OgaRequest* r OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestGetUnseenToken(OgaRequest* request, int32_t* out); /** - * \brief Checks if the request is done processing. + * \brief Checks if the current generation turn is complete. * - * This function checks if the request has finished processing. The request is done when one of the termination - * conditions has been reached (e.g. end of sequence token is encountered or the request was cancelled). - * If the request is done, it will return true; otherwise, it will return false. + * This function returns true at OgaRequestStatus_turn_complete. It does not mean that the request is permanently + * closed; OgaRequestContinue may queue another turn while state remains resident. * * \param[in] request The request to check if it is done. * \param[out] out Boolean flag that will be set to true if the request is done, or false otherwise. @@ -1312,6 +1333,15 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestGetUnseenToken(OgaRequest* request, */ OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestIsDone(const OgaRequest* request, bool* out); +/** + * \brief Gets the request lifecycle status. + * + * \param[in] request The request to inspect. + * \param[out] out The current lifecycle status. + * \return OgaResult containing the error message if the status could not be read, or nullptr on success. + */ +OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestGetStatus(const OgaRequest* request, OgaRequestStatus* out); + /** * \brief Registers an execution provider library with ONNXRuntime API. * \param registration_name name for registration. diff --git a/src/python/python.cpp b/src/python/python.cpp index 0067f2f58b..c7cfe2f5d3 100644 --- a/src/python/python.cpp +++ b/src/python/python.cpp @@ -710,6 +710,13 @@ PYBIND11_MODULE(onnxruntime_genai, m) { .def("unload", &OgaAdapters::UnloadAdapter) .def("load", &OgaAdapters::LoadAdapter); + pybind11::enum_(m, "RequestStatus") + .value("CREATED", OgaRequestStatus_created) + .value("QUEUED", OgaRequestStatus_queued) + .value("IN_PROGRESS", OgaRequestStatus_in_progress) + .value("TURN_COMPLETE", OgaRequestStatus_turn_complete) + .value("CLOSED", OgaRequestStatus_closed); + pybind11::class_(m, "Request") .def(pybind11::init( [](PyGeneratorParams& params) { @@ -721,8 +728,15 @@ PYBIND11_MODULE(onnxruntime_genai, m) { sequences->Append(tokens_span.data(), tokens_span.size()); request.AddTokens(*sequences); }) + .def("continue_with", [](OgaRequest& request, pybind11::array_t tokens) { + auto sequences = OgaSequences::Create(); + auto tokens_span = ToSpan(tokens); + sequences->Append(tokens_span.data(), tokens_span.size()); + request.Continue(*sequences); + }) .def("has_unseen_tokens", &OgaRequest::HasUnseenTokens) .def("is_done", &OgaRequest::IsDone) + .def_property_readonly("status", &OgaRequest::GetStatus) .def("get_unseen_token", &OgaRequest::GetUnseenToken) .def("set_opaque_data", [](OgaRequest& request, pybind11::object opaque_data) { request.SetOpaqueData(opaque_data.ptr()); diff --git a/test/engine/engine_invariants_tests.cpp b/test/engine/engine_invariants_tests.cpp index a96be77567..e99fb9fbd4 100644 --- a/test/engine/engine_invariants_tests.cpp +++ b/test/engine/engine_invariants_tests.cpp @@ -311,14 +311,14 @@ TEST(InvariantValidatorTest, SeenBeyondCurrentReported) { TEST(InvariantValidatorTest, CompletedRequestWithFinalUnprocessedTokenIsValid) { // At completion the just-generated final token is appended but never fed back to the model, so a - // Completed Request legitimately reports one (or more) unprocessed token(s). This must not fire. - auto request = MakeValidRequest(kRequestA, RequestStatus::Completed, 10, 9, 10); + // A TurnComplete Request may legitimately report unprocessed tokens. This must not fire. + auto request = MakeValidRequest(kRequestA, RequestStatus::TurnComplete, 10, 9, 10); EXPECT_TRUE(ValidateRequestInvariants(request).empty()); } TEST(InvariantValidatorTest, CompletedRequestFullyProcessedIsValid) { EXPECT_TRUE(ValidateRequestInvariants( - MakeValidRequest(kRequestA, RequestStatus::Completed, 10, 10, 10)) + MakeValidRequest(kRequestA, RequestStatus::TurnComplete, 10, 10, 10)) .empty()); } diff --git a/test/engine/engine_step_tests.cpp b/test/engine/engine_step_tests.cpp index a561d0b796..97465e3ece 100644 --- a/test/engine/engine_step_tests.cpp +++ b/test/engine/engine_step_tests.cpp @@ -103,6 +103,13 @@ TEST_F(EngineStepTest, FittingRequestsShareOneDecodeAndDrainWithoutReexecuting) // Under capacity backpressure Step decodes only the requests that fit, then forms a fresh batch for // the deferred request on a later step -- one decode per internal cycle, never an over-capacity run. +// +// Cache residency is released only by an explicit RemoveRequest, so a finished request keeps its +// slot until the caller gives it back. With capacity for two requests, the third is admitted only +// after a completed request has been removed. The loop therefore removes each ready request as the +// caller is expected to, which is what frees the slot the deferred request needs; the assertions +// pin both halves of that contract -- the slot is still held when the request is handed back, and +// it is released exactly at removal. TEST_F(EngineStepTest, BackpressureFormsAFreshBatchAcrossSteps) { auto engine = MakeDoublesEngine(model_, /*capacity=*/2, EosToken(*model_)); @@ -115,8 +122,17 @@ TEST_F(EngineStepTest, BackpressureFormsAFreshBatchAcrossSteps) { } std::vector> returned; + int removals = 0; while (auto ready = engine.engine->Step()) { + // The fixture's executor forces end-of-stream, so every request Step hands back is finished. + EXPECT_TRUE(ready->IsDone()); + // The finished request still owns its cache slot: nothing is reclaimed implicitly. + EXPECT_EQ(engine.cache->deallocate_calls, removals); returned.push_back(ready); + + engine.engine->RemoveRequest(ready); + ++removals; + EXPECT_EQ(engine.cache->deallocate_calls, removals); } EXPECT_EQ(returned.size(), 3u); @@ -125,8 +141,9 @@ TEST_F(EngineStepTest, BackpressureFormsAFreshBatchAcrossSteps) { std::sort(returned.begin(), returned.end()); EXPECT_EQ(returned, sorted_requests); for (const auto& request : requests) { - EXPECT_TRUE(request->IsDone()); + EXPECT_EQ(request->status_, RequestStatus::Closed); } + EXPECT_EQ(engine.cache->AllocatedCount(), 0u); ASSERT_EQ(engine.executor->decoded_batch_sizes.size(), 2u); EXPECT_EQ(engine.executor->decoded_batch_sizes[0], 2u); EXPECT_EQ(engine.executor->decoded_batch_sizes[1], 1u); @@ -149,6 +166,7 @@ TEST_F(EngineStepTest, StaticBatchingRetainsLegacyCommitOrdering) { auto scheduler = Scheduler::Create(model_, cache); auto executor = std::make_unique( model_, cache, EosToken(*model_), trace); + auto* cache_observer = cache.get(); auto* executor_observer = executor.get(); EngineDependencies dependencies{cache, std::move(scheduler), std::move(executor)}; @@ -160,6 +178,95 @@ TEST_F(EngineStepTest, StaticBatchingRetainsLegacyCommitOrdering) { EXPECT_EQ(engine->Step(), request); EXPECT_EQ(executor_observer->decode_calls, 1); EXPECT_LT(IndexOf(*trace, "Allocate"), IndexOf(*trace, "Decode")); + + const int allocations_before = cache_observer->allocate_calls; + const std::vector continuation{5, 6}; + request->Continue(continuation); + ASSERT_EQ(request->status_, RequestStatus::Assigned); + + EXPECT_EQ(engine->Step(), request); + EXPECT_EQ(request->status_, RequestStatus::TurnComplete); + EXPECT_EQ(executor_observer->decode_calls, 2); + EXPECT_EQ(cache_observer->allocate_calls, allocations_before); +} + +TEST_F(EngineStepTest, StaticAdmissionValidationDoesNotStrandRequest) { + model_->config_->engine.dynamic_batching.reset(); + auto cache = std::make_shared( + model_, /*capacity=*/4, nullptr, /*supports_dynamic_batching=*/false); + auto scheduler = Scheduler::Create(model_, cache); + auto executor = std::make_unique( + model_, cache, EosToken(*model_)); + EngineDependencies dependencies{cache, std::move(scheduler), + std::move(executor)}; + auto engine = std::make_shared(model_, std::move(dependencies)); + auto prompt = Prompt(10); + auto request = MintRequest(*model_, prompt); + request->Params()->search.chunk_size = 2; + + EXPECT_THROW(engine->AddRequest(request), std::runtime_error); + EXPECT_EQ(request->status_, RequestStatus::Unassigned); + + request->Params()->search.chunk_size.reset(); + EXPECT_NO_THROW(engine->AddRequest(request)); + EXPECT_EQ(request->status_, RequestStatus::Assigned); +} + +TEST_F(EngineStepTest, StaticContinueFailsAfterBatchRecycling) { + model_->config_->engine.dynamic_batching.reset(); + auto cache = std::make_shared(model_); + auto scheduler = Scheduler::Create(model_, cache); + auto executor = std::make_unique( + model_, cache, EosToken(*model_)); + EngineDependencies dependencies{cache, std::move(scheduler), + std::move(executor)}; + auto engine = std::make_shared(model_, std::move(dependencies)); + + auto first_prompt = Prompt(10); + auto first = MintRequest(*model_, first_prompt); + engine->AddRequest(first); + ASSERT_EQ(engine->Step(), first); + ASSERT_EQ(first->status_, RequestStatus::TurnComplete); + + auto second_prompt = Prompt(20); + auto second = MintRequest(*model_, second_prompt); + engine->AddRequest(second); + for (int step = 0; step < 2 && cache->IsResident(first); ++step) { + ASSERT_NE(engine->Step(), nullptr); + } + ASSERT_FALSE(cache->IsResident(first)); + + const std::vector continuation{5, 6}; + EXPECT_THROW(first->Continue(continuation), std::runtime_error); + EXPECT_EQ(first->status_, RequestStatus::TurnComplete); +} + +TEST_F(EngineStepTest, StaticContinueRejectsMultiRowBatchAfterPeerCloses) { + model_->config_->engine.dynamic_batching.reset(); + auto cache = std::make_shared( + model_, /*capacity=*/4, nullptr, /*supports_dynamic_batching=*/false); + auto scheduler = Scheduler::Create(model_, cache); + auto executor = std::make_unique( + model_, cache, EosToken(*model_)); + EngineDependencies dependencies{cache, std::move(scheduler), + std::move(executor)}; + auto engine = std::make_shared(model_, std::move(dependencies)); + + auto first_prompt = Prompt(10); + auto second_prompt = Prompt(20); + auto first = MintRequest(*model_, first_prompt); + auto second = MintRequest(*model_, second_prompt); + engine->AddRequest(first); + engine->AddRequest(second); + ASSERT_EQ(engine->Step(), first); + ASSERT_EQ(first->status_, RequestStatus::TurnComplete); + ASSERT_EQ(second->status_, RequestStatus::TurnComplete); + + engine->RemoveRequest(second); + ASSERT_EQ(second->status_, RequestStatus::Closed); + const std::vector continuation{5, 6}; + EXPECT_THROW(first->Continue(continuation), std::runtime_error); + EXPECT_EQ(first->status_, RequestStatus::TurnComplete); } TEST_F(EngineStepTest, StepDoesNotReturnNullWhenCapacityDefersPendingWork) { @@ -207,6 +314,89 @@ TEST_F(EngineStepTest, RetryableExecutionFailureRollsBackAndCanRetry) { EXPECT_TRUE(request->IsDone()); } +TEST_F(EngineStepTest, ContinuedResidentRollsBackToQueuedAndCanRetry) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); + auto prompt = Prompt(10); + auto request = MintRequest(*model_, prompt); + engine.engine->AddRequest(request); + ASSERT_EQ(engine.engine->Step(), request); + ASSERT_EQ(request->status_, RequestStatus::TurnComplete); + ASSERT_EQ(engine.cache->AllocatedCount(), 1u); + + const std::vector continuation{5, 6}; + request->Continue(continuation); + const auto before = request->Snapshot(); + ASSERT_EQ(before.status, RequestStatus::Assigned); + engine.executor->SetNextFailure( + ScriptedExecutionFailure::RetryableDuringExecution); + + try { + static_cast(engine.engine->Step()); + FAIL() << "Expected retryable execution failure."; + } catch (const EngineStepError& error) { + EXPECT_EQ(error.Outcome().kind, StepOutcomeKind::RetryableBatchAbort); + } + + const auto rolled_back = request->Snapshot(); + EXPECT_EQ(rolled_back.status, RequestStatus::Assigned); + EXPECT_EQ(rolled_back.current_sequence_length, before.current_sequence_length); + EXPECT_EQ(rolled_back.processed_sequence_length, + before.processed_sequence_length); + EXPECT_EQ(engine.cache->AllocatedCount(), 1u); + + auto ready = engine.engine->Step(); + EXPECT_EQ(ready, request); + EXPECT_EQ(request->status_, RequestStatus::TurnComplete); +} + +TEST_F(EngineStepTest, ContinuedResidentUsesChunkedPrefillBeforeSampling) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); + auto prompt = Prompt(10); + auto request = MintRequest(*model_, prompt); + request->Params()->search.chunk_size = 2; + engine.engine->AddRequest(request); + ASSERT_EQ(engine.engine->Step(), request); + ASSERT_EQ(request->status_, RequestStatus::TurnComplete); + + const size_t calls_before = engine.executor->decoded_token_counts.size(); + const std::vector continuation{5, 6, 7, 8, 9}; + request->Continue(continuation); + ASSERT_EQ(engine.engine->Step(), request); + + ASSERT_EQ(engine.executor->decoded_token_counts.size(), calls_before + 3); + EXPECT_EQ(engine.executor->decoded_token_counts[calls_before], 2u); + EXPECT_EQ(engine.executor->decoded_token_counts[calls_before + 1], 2u); + EXPECT_EQ(engine.executor->decoded_token_counts[calls_before + 2], 1u); + EXPECT_EQ(request->ProcessedSequenceLength(), + request->CurrentSequenceLength()); + EXPECT_EQ(request->status_, RequestStatus::TurnComplete); +} + +TEST_F(EngineStepTest, UnserviceableContinuationRemainsQueuedUntilClosed) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); + auto prompt = Prompt(10); + auto request = MintRequest(*model_, prompt); + engine.engine->AddRequest(request); + ASSERT_EQ(engine.engine->Step(), request); + ASSERT_EQ(request->status_, RequestStatus::TurnComplete); + + const std::vector continuation{5, 6}; + request->Continue(continuation); + engine.cache->SetUnserviceableRequest(request); + + try { + static_cast(engine.engine->Step()); + FAIL() << "Expected an unserviceable continuation error."; + } catch (const EngineStepError& error) { + EXPECT_EQ(error.Outcome().kind, StepOutcomeKind::UnserviceableRequest); + } + + EXPECT_EQ(request->status_, RequestStatus::Assigned); + EXPECT_EQ(engine.cache->AllocatedCount(), 1u); + EXPECT_NO_THROW(engine.engine->RemoveRequest(request)); + EXPECT_EQ(request->status_, RequestStatus::Closed); +} + TEST_F(EngineStepTest, ExecutionCapacityFailureRollsBackWithoutPoisoningEngine) { auto engine = MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); auto prompt = Prompt(10); @@ -309,6 +499,95 @@ TEST_F(EngineStepTest, LaterRequestFailureRestoresEarlierSample) { second_before.current_sequence_length + 1); } +TEST_F(EngineStepTest, RemovingUndrainedReadyRequestPurgesItFromQueue) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/8, /*forced_token=*/5); + auto first_prompt = Prompt(10); + auto second_prompt = Prompt(20); + auto first = MintRequest(*model_, first_prompt); + auto second = MintRequest(*model_, second_prompt); + engine.engine->AddRequest(first); + engine.engine->AddRequest(second); + + ASSERT_EQ(engine.engine->Step(), first); + ASSERT_TRUE(second->HasUnseenTokens()); + engine.engine->RemoveRequest(second); + ASSERT_EQ(second->status_, RequestStatus::Closed); + + EXPECT_EQ(engine.engine->Step(), first); +} + +TEST_F(EngineStepTest, RemovingOnlyDrainedReadyRequestClearsQueue) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/8, /*forced_token=*/5); + auto prompt = Prompt(10); + auto request = MintRequest(*model_, prompt); + engine.engine->AddRequest(request); + + ASSERT_EQ(engine.engine->Step(), request); + engine.engine->RemoveRequest(request); + + EXPECT_EQ(request->status_, RequestStatus::Closed); + EXPECT_EQ(engine.engine->Step(), nullptr); +} + +TEST_F(EngineStepTest, RepeatedReadyRemovalsPreserveRemainingQueueOrder) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/8, /*forced_token=*/5); + std::vector> requests; + for (int32_t seed : {10, 20, 30}) { + auto prompt = Prompt(seed); + auto request = MintRequest(*model_, prompt); + engine.engine->AddRequest(request); + requests.push_back(std::move(request)); + } + + ASSERT_EQ(engine.engine->Step(), requests[0]); + engine.engine->RemoveRequest(requests[0]); + // The first removal resets the drain cursor to zero while two ready entries remain. Removing a + // second request must compact that nonempty queue in place without overlapping std::move ranges. + engine.engine->RemoveRequest(requests[1]); + + EXPECT_EQ(requests[0]->status_, RequestStatus::Closed); + EXPECT_EQ(requests[1]->status_, RequestStatus::Closed); + EXPECT_EQ(engine.engine->Step(), requests[2]); +} + +TEST_F(EngineStepTest, StaticTurnCompleteRowIsNotRepublishedWhilePeerRuns) { + model_->config_->engine.dynamic_batching.reset(); + auto cache = std::make_shared( + model_, /*capacity=*/4, nullptr, /*supports_dynamic_batching=*/false); + auto scheduler = Scheduler::Create(model_, cache); + auto executor = std::make_unique( + model_, cache, /*forced_token=*/5); + auto* executor_observer = executor.get(); + EngineDependencies dependencies{cache, std::move(scheduler), + std::move(executor)}; + auto engine = std::make_shared(model_, std::move(dependencies)); + + auto first_prompt = Prompt(10); + auto second_prompt = Prompt(20); + auto first = MintRequest(*model_, first_prompt); + auto second = MintRequest(*model_, second_prompt); + first->Params()->search.max_length = + static_cast(first_prompt.size() + 1); + second->Params()->search.max_length = + static_cast(second_prompt.size() + 3); + engine->AddRequest(first); + engine->AddRequest(second); + + ASSERT_EQ(engine->Step(), first); + ASSERT_EQ(first->status_, RequestStatus::TurnComplete); + while (first->HasUnseenTokens()) { + static_cast(first->UnseenToken()); + } + ASSERT_EQ(engine->Step(), second); // Drain the other result from the same model run. + while (second->HasUnseenTokens()) { + static_cast(second->UnseenToken()); + } + ASSERT_EQ(executor_observer->decode_calls, 1); + + EXPECT_EQ(engine->Step(), second); + EXPECT_EQ(executor_observer->decode_calls, 2); +} + TEST_F(EngineStepTest, FatalExecutionFailureMarksEngineUnhealthy) { auto engine = MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); auto prompt = Prompt(10); @@ -329,7 +608,29 @@ TEST_F(EngineStepTest, FatalExecutionFailureMarksEngineUnhealthy) { EXPECT_EQ(engine.executor->decode_calls, 1); EXPECT_EQ(engine.cache->AllocatedCount(), 0u); request->Remove(); - EXPECT_EQ(request->status_, RequestStatus::Unassigned); + EXPECT_EQ(request->status_, RequestStatus::Closed); +} + +TEST_F(EngineStepTest, ContinueIsRejectedAfterEngineBecomesUnhealthy) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); + auto first_prompt = Prompt(10); + auto first = MintRequest(*model_, first_prompt); + engine.engine->AddRequest(first); + ASSERT_EQ(engine.engine->Step(), first); + ASSERT_EQ(first->status_, RequestStatus::TurnComplete); + + auto second_prompt = Prompt(20); + auto second = MintRequest(*model_, second_prompt); + engine.engine->AddRequest(second); + engine.executor->SetNextFailure(ScriptedExecutionFailure::Fatal); + EXPECT_THROW(static_cast(engine.engine->Step()), EngineStepError); + + const auto before = first->Snapshot(); + const std::vector continuation{5, 6}; + EXPECT_THROW(first->Continue(continuation), EngineStepError); + const auto after = first->Snapshot(); + EXPECT_EQ(after.status, RequestStatus::TurnComplete); + EXPECT_EQ(after.current_sequence_length, before.current_sequence_length); } TEST_F(EngineStepTest, UnserviceableRequestDoesNotBlockFittingRequest) { diff --git a/test/engine/engine_test_doubles.h b/test/engine/engine_test_doubles.h index 53727a9df1..aea4676fc7 100644 --- a/test/engine/engine_test_doubles.h +++ b/test/engine/engine_test_doubles.h @@ -81,6 +81,12 @@ struct RecordingCacheManager : CacheManager { std::vector> AllocatedRequests() const override { return allocated_; } + bool IsResident(const std::shared_ptr& request) const override { + return std::find(allocated_.begin(), allocated_.end(), request) != allocated_.end(); + } + + size_t ResidentRequestCount() const override { return allocated_.size(); } + StepPlanningResult PlanStepResources(StepPlan& plan) const override { const size_t request_limit = plan.scheduled_request_limit == 0 ? capacity_ : plan.scheduled_request_limit; diff --git a/test/engine/request_lifecycle_tests.cpp b/test/engine/request_lifecycle_tests.cpp index 52992bbf10..2e2b38bf61 100644 --- a/test/engine/request_lifecycle_tests.cpp +++ b/test/engine/request_lifecycle_tests.cpp @@ -4,8 +4,8 @@ // Lifecycle tests for the engine Request state machine. Because a tiny real // CPU fixture model is available, these tests drive genuine Request objects (rather than a mock // Search) and pin the transition policy: which mutations each status permits, and how -// create/assign/schedule/remove move a request between Unassigned, Assigned, InProgress, and -// Completed. +// create/assign/schedule/continue/remove move a request between Unassigned, Assigned, InProgress, +// TurnComplete, and Closed. #include #include @@ -48,10 +48,18 @@ class RequestLifecycleTest : public ::testing::Test { DoublesEngine engine_; }; +TEST(ContinuousDecodingDeviceSupportTest, MatchesKvCacheCapabilityContract) { + EXPECT_TRUE(SupportsContinuousDecoding(DeviceType::CPU)); + EXPECT_TRUE(SupportsContinuousDecoding(DeviceType::CUDA)); + EXPECT_FALSE(SupportsContinuousDecoding(DeviceType::DML)); + EXPECT_FALSE(SupportsContinuousDecoding(DeviceType::QnnHtp)); +} + // A request must carry at least one token per append; an empty batch is rejected. TEST_F(RequestLifecycleTest, EmptyAppendIsRejected) { auto request = NewRequest(); EXPECT_THROW(request->AddTokens({}), std::runtime_error); + EXPECT_THROW(request->Continue({}), std::runtime_error); } TEST_F(RequestLifecycleTest, EmptyRequestIsRejectedBeforeAssignment) { @@ -122,31 +130,145 @@ TEST_F(RequestLifecycleTest, AppendIsRejectedWhileInProgress) { EXPECT_EQ(request->CurrentSequenceLength(), length_before); } -// After a request completes, appending tokens resumes its sequence (the continuation path) rather -// than being rejected. -TEST_F(RequestLifecycleTest, AppendAfterCompletedExtendsSequence) { +TEST_F(RequestLifecycleTest, AddTokensIsRejectedAfterSubmission) { auto prompt = Prompt(); + const std::vector more{5}; auto request = MintAssignedRequest(engine_.engine, *model_, prompt); - const int64_t assigned_length = request->CurrentSequenceLength(); - request->status_ = RequestStatus::Completed; + EXPECT_THROW(request->AddTokens(more), std::runtime_error); + request->Schedule(); + EXPECT_THROW(request->AddTokens(more), std::runtime_error); +} + +// After a turn completes, Continue appends another input fragment and queues the resident request. +TEST_F(RequestLifecycleTest, ContinueAfterTurnCompleteQueuesNextTurn) { + auto prompt = Prompt(); + auto request = MintRequest(*model_, prompt); + const int64_t assigned_length = static_cast(prompt.size()); + engine_.engine->AddRequest(request); + + engine_.engine->Step(); + ASSERT_EQ(request->status_, RequestStatus::TurnComplete); std::vector more{5, 6}; - request->AddTokens(more); + EXPECT_THROW(request->AddTokens(more), std::runtime_error); + request->Continue(more); EXPECT_EQ(request->CurrentSequenceLength(), assigned_length + static_cast(more.size())); + EXPECT_EQ(request->status_, RequestStatus::Assigned); + EXPECT_FALSE(request->IsDone()); + EXPECT_EQ(engine_.engine->Step(), request); + EXPECT_EQ(request->status_, RequestStatus::TurnComplete); +} + +TEST_F(RequestLifecycleTest, ContinueBeyondContextIsRejectedBeforeMutation) { + auto prompt = Prompt(); + auto request = MintRequest(*model_, prompt); + engine_.engine->AddRequest(request); + engine_.engine->Step(); + ASSERT_EQ(request->status_, RequestStatus::TurnComplete); + const auto before = request->Snapshot(); + + const size_t remaining = + static_cast(request->Params()->search.max_length - + request->CurrentSequenceLength()); + std::vector too_many(remaining, 5); + EXPECT_THROW(request->Continue(too_many), std::runtime_error); + + const auto after = request->Snapshot(); + EXPECT_EQ(after.status, before.status); + EXPECT_EQ(after.current_sequence_length, before.current_sequence_length); + EXPECT_EQ(after.processed_sequence_length, before.processed_sequence_length); } -// Removing a request releases it from the engine (deallocating its cache resources) and returns it -// to the Unassigned state. -TEST_F(RequestLifecycleTest, RemoveReturnsRequestToUnassigned) { +TEST_F(RequestLifecycleTest, ContinuePreservesUnreadOutputAndHidesInputTokens) { auto prompt = Prompt(); auto request = MintAssignedRequest(engine_.engine, *model_, prompt); + engine_.cache->Allocate({request}); + request->Schedule(); + + RequestStepPlan first_plan; + first_plan.request = request; + first_plan.request_id = request.get(); + first_plan.sequence_length_before = request->CurrentSequenceLength(); + first_plan.target_cache_slots = + static_cast(first_plan.sequence_length_before); + constexpr int32_t generated_token = 5; + auto first_logits = LogitsForToken(*model_, generated_token); + request->SaveStateForTransaction(); + const auto first_result = request->ApplyLogitsForTransaction(first_logits); + request->CommitStateForTransaction(); + request->CommitStep(first_plan, first_result); + ASSERT_TRUE(request->HasUnseenTokens()); + + RequestStepPlan completion_plan; + completion_plan.request = request; + completion_plan.request_id = request.get(); + completion_plan.sequence_length_before = request->CurrentSequenceLength(); + completion_plan.target_cache_slots = + static_cast(completion_plan.sequence_length_before); + auto eos_logits = LogitsForToken(*model_, EosToken(*model_)); + request->SaveStateForTransaction(); + const auto completion_result = + request->ApplyLogitsForTransaction(eos_logits); + request->CommitStateForTransaction(); + request->CommitStep(completion_plan, completion_result); + ASSERT_EQ(request->status_, RequestStatus::TurnComplete); + + const std::vector continuation{6, 7}; + request->Continue(continuation); + + EXPECT_EQ(request->status_, RequestStatus::Assigned); + ASSERT_TRUE(request->HasUnseenTokens()); + EXPECT_EQ(request->UnseenToken(), generated_token); + EXPECT_FALSE(request->HasUnseenTokens()); +} + +TEST_F(RequestLifecycleTest, ContinueIsRejectedOutsideTurnComplete) { + const std::vector more{5}; + auto request = NewRequest(); + EXPECT_THROW(request->Continue(more), std::runtime_error); + + auto prompt = Prompt(); + request->AddTokens(prompt); + engine_.engine->AddRequest(request); + EXPECT_THROW(request->Continue(more), std::runtime_error); + + request->Schedule(); + EXPECT_THROW(request->Continue(more), std::runtime_error); +} + +// Removing a request releases it from the engine and makes it terminal. +TEST_F(RequestLifecycleTest, RemoveMakesRequestTerminal) { + auto prompt = Prompt(); + const std::vector more{5}; + auto request = MintAssignedRequest(engine_.engine, *model_, prompt); ASSERT_EQ(request->status_, RequestStatus::Assigned); request->Remove(); - EXPECT_EQ(request->status_, RequestStatus::Unassigned); + EXPECT_EQ(request->status_, RequestStatus::Closed); EXPECT_EQ(engine_.cache->deallocate_calls, 1); EXPECT_EQ(engine_.cache->AllocatedCount(), 0u); + EXPECT_THROW(request->AddTokens(more), std::runtime_error); + EXPECT_THROW(request->Continue(more), std::runtime_error); + EXPECT_THROW(request->Remove(), std::runtime_error); +} + +TEST_F(RequestLifecycleTest, RemoveIsRejectedBeforeSubmission) { + auto request = NewRequest(); + EXPECT_THROW(request->Remove(), std::runtime_error); + EXPECT_EQ(request->status_, RequestStatus::Unassigned); +} + +TEST_F(RequestLifecycleTest, RemoveClosesRequestAfterEngineDestruction) { + auto local_engine = + MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); + auto prompt = Prompt(); + auto request = MintRequest(*model_, prompt); + local_engine.engine->AddRequest(request); + local_engine.engine.reset(); + + EXPECT_NO_THROW(request->Remove()); + EXPECT_EQ(request->status_, RequestStatus::Closed); } TEST_F(RequestLifecycleTest, TransactionalLogitsStageUntilCommit) { @@ -225,7 +347,7 @@ TEST_F(RequestLifecycleTest, PartialPrefillAdvancesOnlyAtCommit) { EXPECT_FALSE(request->HasUnseenTokens()); } -TEST_F(RequestLifecycleTest, FirstTransactionalStepCanCommitDirectlyToCompleted) { +TEST_F(RequestLifecycleTest, FirstTransactionalStepCanCommitDirectlyToTurnComplete) { auto prompt = Prompt(); auto request = MintAssignedRequest(engine_.engine, *model_, prompt); const auto before = request->Snapshot(); @@ -242,7 +364,7 @@ TEST_F(RequestLifecycleTest, FirstTransactionalStepCanCommitDirectlyToCompleted) request->CommitStep(plan, result); EXPECT_TRUE(result.done); - EXPECT_EQ(request->status_, RequestStatus::Completed); + EXPECT_EQ(request->status_, RequestStatus::TurnComplete); } TEST_F(RequestLifecycleTest, RequestRejectsMultiSequenceSearch) { diff --git a/test/engine/scheduler_contract_tests.cpp b/test/engine/scheduler_contract_tests.cpp index 790a2a37b0..66b4536c3e 100644 --- a/test/engine/scheduler_contract_tests.cpp +++ b/test/engine/scheduler_contract_tests.cpp @@ -145,7 +145,7 @@ TEST_F(SchedulerContractTest, DynamicHonorsCapacityBackpressure) { EXPECT_EQ(third->status_, RequestStatus::Assigned); } -TEST_F(SchedulerContractTest, DynamicDeallocatesCompletedRequests) { +TEST_F(SchedulerContractTest, DynamicRetainsTurnCompleteRequestsUntilExplicitRemoval) { auto cache = std::make_shared(model_, /*capacity=*/8); DynamicBatchScheduler scheduler(model_, cache); @@ -153,7 +153,7 @@ TEST_F(SchedulerContractTest, DynamicDeallocatesCompletedRequests) { MakePrefillResident(scheduler, *cache, request); ASSERT_EQ(cache->AllocatedCount(), 1u); - request->status_ = RequestStatus::Completed; + request->status_ = RequestStatus::TurnComplete; auto second = Assigned(20); scheduler.AddRequest(second); StepPlan plan; @@ -161,10 +161,84 @@ TEST_F(SchedulerContractTest, DynamicDeallocatesCompletedRequests) { const auto result = scheduler.PlanStep(plan); ASSERT_TRUE(result.executable); - EXPECT_GE(cache->deallocate_calls, 1); - EXPECT_EQ(cache->AllocatedCount(), 0u); + EXPECT_EQ(cache->deallocate_calls, 0); + EXPECT_EQ(cache->AllocatedCount(), 1u); ASSERT_EQ(plan.requests.size(), 1u); EXPECT_EQ(plan.requests[0].request, second); + + scheduler.RemoveRequest(request); + + EXPECT_EQ(cache->deallocate_calls, 1); + EXPECT_EQ(cache->AllocatedCount(), 0u); + EXPECT_FALSE(cache->IsResident(request)); +} + +TEST_F(SchedulerContractTest, DynamicTurnCompleteResidencyAppliesCapacityBackpressure) { + auto cache = std::make_shared(model_, /*capacity=*/1); + DynamicBatchScheduler scheduler(model_, cache); + + auto completed = Assigned(10); + MakePrefillResident(scheduler, *cache, completed); + completed->status_ = RequestStatus::TurnComplete; + + auto waiting = Assigned(20); + scheduler.AddRequest(waiting); + StepPlan plan; + + const auto blocked = scheduler.PlanStep(plan); + + EXPECT_FALSE(blocked.executable); + EXPECT_TRUE(blocked.capacity_deferred); + EXPECT_EQ(cache->deallocate_calls, 0); + EXPECT_EQ(cache->AllocatedCount(), 1u); + EXPECT_TRUE(plan.requests.empty()); + + scheduler.RemoveRequest(completed); + const auto admitted = scheduler.PlanStep(plan); + + ASSERT_TRUE(admitted.executable); + EXPECT_EQ(cache->deallocate_calls, 1); + EXPECT_EQ(cache->AllocatedCount(), 0u); + ASSERT_EQ(plan.requests.size(), 1u); + EXPECT_EQ(plan.requests[0].request, waiting); +} + +TEST_F(SchedulerContractTest, DynamicResidentQueuedRequestIsNotReadmitted) { + auto cache = std::make_shared(model_, /*capacity=*/8); + DynamicBatchScheduler scheduler(model_, cache); + + auto request = Assigned(10); + MakeDecodeResident(scheduler, *cache, request); + request->status_ = RequestStatus::Assigned; + StepPlan plan; + + const auto result = scheduler.PlanStep(plan); + + ASSERT_TRUE(result.executable); + ASSERT_EQ(plan.requests.size(), 1u); + EXPECT_EQ(plan.requests[0].request, request); + EXPECT_FALSE(plan.requests[0].newly_admitted); + EXPECT_EQ(cache->AllocatedCount(), 1u); +} + +TEST_F(SchedulerContractTest, StaticResidentQueuedRequestSchedulesWithoutAllocation) { + model_->config_->engine.dynamic_batching.reset(); + auto cache = std::make_shared( + model_, /*capacity=*/4, nullptr, /*supports_dynamic_batching=*/false); + StaticBatchScheduler scheduler(model_, cache); + + auto request = Assigned(10); + scheduler.AddRequest(request); + cache->Allocate({request}); + request->status_ = RequestStatus::Assigned; + const int allocations_before = cache->allocate_calls; + + auto scheduled = scheduler.Schedule(); + + ASSERT_EQ(scheduled.size(), 1u); + EXPECT_EQ(scheduled[0], request); + EXPECT_EQ(request->status_, RequestStatus::InProgress); + EXPECT_EQ(cache->allocate_calls, allocations_before); } // Removing a request from the dynamic scheduler deallocates its cache resources immediately. diff --git a/test/python/test_onnxruntime_genai_engine.py b/test/python/test_onnxruntime_genai_engine.py index 594697361d..999d83b0e5 100644 --- a/test/python/test_onnxruntime_genai_engine.py +++ b/test/python/test_onnxruntime_genai_engine.py @@ -208,6 +208,122 @@ def test_completion_isolation(model): assert long_sink.tokens == long_isolated, "survivor diverged after its sibling completed" +def test_continuation_while_peer_remains_active(model): + short_max_new, long_max_new = 60, 80 + # EOS is valid input context here; Continue must reset the prior turn's done state rather than + # treating an EOS token in the new prompt fragment as a newly generated stop. + follow_up = [_EOS_TOKEN_ID, 12] + + reference_engine = og.Engine(model) + reference_sink = _Sink() + reference = _add_request( + reference_engine, model, _PROMPT_A, short_max_new, reference_sink + ) + while not reference.is_done(): + ready = reference_engine.step() + assert ready is not None + _drain(ready) + reference.continue_with(np.asarray(follow_up, dtype=np.int32)) + _run(reference_engine) + + engine = og.Engine(model) + short_sink, long_sink = _Sink(), _Sink() + short = _add_request(engine, model, _PROMPT_A, short_max_new, short_sink) + long = _add_request(engine, model, _PROMPT_LONG, long_max_new, long_sink) + + while not short.is_done(): + ready = engine.step() + assert ready is not None + if _drain(ready) and ready is not short: + engine.remove_request(ready) + + assert not long.is_done(), "peer must remain active when continuation is appended" + for _ in range(3): + ready = engine.step() + assert ready is not None + _drain(ready) + assert not long.is_done(), "peer must remain active during the continuation delay" + + short.continue_with(np.asarray(follow_up, dtype=np.int32)) + _run(engine) + + assert short_sink.tokens == reference_sink.tokens + + +def test_request_rejects_empty_input(model): + params = og.GeneratorParams(model) + request = og.Request(params) + + with pytest.raises(RuntimeError, match="at least one token"): + request.add_tokens(np.asarray([], dtype=np.int32)) + + +def test_request_rejects_tokens_while_awaiting_admission(model): + engine = og.Engine(model) + sink = _Sink() + request = _add_request(engine, model, _PROMPT_A, 8, sink) + + with pytest.raises(RuntimeError, match="initial input before submission"): + request.add_tokens(np.asarray([12], dtype=np.int32)) + + engine.remove_request(request) + + +def test_closed_request_cannot_continue(model): + engine = og.Engine(model) + sink = _Sink() + request = _add_request(engine, model, _PROMPT_A, 8, sink) + engine.remove_request(request) + + assert request.status == og.RequestStatus.CLOSED + with pytest.raises(RuntimeError, match="closed request"): + request.add_tokens(np.asarray([12], dtype=np.int32)) + with pytest.raises(RuntimeError, match="closed request"): + request.continue_with(np.asarray([12], dtype=np.int32)) + + +def test_request_cannot_be_removed_from_another_engine(model): + owner = og.Engine(model) + other = og.Engine(model) + sink = _Sink() + request = _add_request(owner, model, _PROMPT_A, 8, sink) + + with pytest.raises(RuntimeError, match="does not belong"): + other.remove_request(request) + + owner.remove_request(request) + + +def test_request_lifecycle_status(model): + params = og.GeneratorParams(model) + params.set_search_options(do_sample=False, max_length=64) + request = og.Request(params) + assert request.status == og.RequestStatus.CREATED + + request.add_tokens(np.asarray(_PROMPT_A, dtype=np.int32)) + sink = _Sink() + request.set_opaque_data(sink) + engine = og.Engine(model) + engine.add_request(request) + assert request.status == og.RequestStatus.QUEUED + + while not request.is_done(): + ready = engine.step() + assert ready is not None + _drain(ready) + assert request.status == og.RequestStatus.TURN_COMPLETE + + with pytest.raises(RuntimeError, match="use Continue"): + request.add_tokens(np.asarray([12], dtype=np.int32)) + request.continue_with(np.asarray([12], dtype=np.int32)) + assert request.status == og.RequestStatus.QUEUED + + engine.remove_request(request) + assert request.status == og.RequestStatus.CLOSED + with pytest.raises(RuntimeError, match="already closed"): + engine.remove_request(request) + + def test_remove_request_freezes_output(model): max_new = 40 sibling_new = 16 From 6211fd32b8f40ef7081c833cafc094304f0e0087 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 15:13:50 -0500 Subject: [PATCH 02/19] Refine continuous decoding lifecycle handling Simplify ready-result removal, reuse cache residency queries, and make terminal continuation errors take precedence. Trim redundant and out-of-scope coverage while preserving continuation, retention, rollback, static-batch, and public lifecycle contracts. Files: docs/paged_attention_engine.md; src/engine/{engine,request,scheduler}.{cpp,h}; test/engine/{engine_invariants,engine_step,request_lifecycle,scheduler_contract}_tests.cpp; test/python/test_onnxruntime_genai_engine.py Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b232f82c-f25c-429a-a855-7f8c8f40bbf3 --- docs/paged_attention_engine.md | 48 ++++++++------------ src/engine/engine.cpp | 14 ++---- src/engine/request.cpp | 8 ++-- src/engine/scheduler.cpp | 19 ++------ src/engine/scheduler.h | 4 -- test/engine/engine_invariants_tests.cpp | 6 +-- test/engine/engine_step_tests.cpp | 45 +----------------- test/engine/request_lifecycle_tests.cpp | 4 +- test/engine/scheduler_contract_tests.cpp | 20 -------- test/python/test_onnxruntime_genai_engine.py | 25 ++-------- 10 files changed, 39 insertions(+), 154 deletions(-) diff --git a/docs/paged_attention_engine.md b/docs/paged_attention_engine.md index ac08d75bcf..81b40fd3d6 100644 --- a/docs/paged_attention_engine.md +++ b/docs/paged_attention_engine.md @@ -137,11 +137,9 @@ the model's chat template. `Continue(tokens)` appends the next input fragment and moves a resident request back to `Assigned`. `AddTokens()` remains an initial-input-only operation. -There is no fixed wall-clock or next-step timeout. 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. Until Phase 2 defines -residency and eviction, insufficient capacity is surfaced as backpressure rather than silently -discarding another conversation's model state. +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. ### `Remove()` @@ -157,10 +155,6 @@ scheduler ownership. Returning to `Unassigned` would imply that the same logical submitted as a new request. A closed static-batch row may remain physically allocated until the batch is recycled, but it is no longer sampled or returned. -Lifecycle status and residency are separate concepts. Phase 1 guarantees that dynamic -`TurnComplete` requests stay resident until `Remove()`. Phase 2 will define observable residency and -automatic eviction; no eviction policy is part of this lifecycle change. - ## The request length counters Three views of request progress are important: @@ -169,7 +163,16 @@ Three views of request progress are important: | --- | --- | | `CurrentSequenceLength()` | Number of tokens currently held by the request's search sequence | | `processed_sequence_length_` | Number of sequence tokens already represented in the committed KV cache | -| `seen_sequence_length_` | High-water sequence index of generated output consumed by the API caller; continuation input may create gaps | +| `seen_sequence_length_` | High-water sequence index of generated output consumed by the API caller; copied into invariant snapshots rather than used to select the next output token | + +Generated-output delivery uses separate bookkeeping because continuation input creates gaps in the +logical sequence: + +| Value | Meaning | +| --- | --- | +| `tokens_host_` | Host-side mirror of the complete logical sequence, including prompt, generated output, and continuation input | +| `unseen_token_indices_` | Positions of generated tokens in `tokens_host_`; continuation-input positions are never added | +| `next_unseen_token_index_` | Cursor into `unseen_token_indices_`; entries at and after this cursor have not been consumed | The unprocessed tokens are: @@ -227,19 +230,16 @@ If the engine has previously encountered a fatal transaction or execution failur `DynamicBatchScheduler::PlanStep()` skips `TurnComplete` residents and builds candidates from executable residents plus waiting requests. -The cache manager checks whether those candidates fit alongside dormant turn-complete requests. It -does not reclaim another request as a side effect of `Step()`. If retained residency prevents -admission or cache growth, the plan reports capacity backpressure; the application decides which -conversation to release with `Remove()`. +The cache manager checks whether those candidates fit alongside dormant turn-complete requests. If +retained residency prevents admission or cache growth, the plan reports capacity backpressure; the +application decides which conversation to release with `Remove()`. ### 2. Build the initial step plan The scheduler snapshots requests that already belong to the paged cache. Executable residents may be `InProgress` or `Assigned`; an `Assigned` resident is a queued continuation. -It then snapshots nonresident waiting requests from the scheduler pool. These are `Assigned` and -are marked as newly admitted candidates. Residency, not status alone, determines -`newly_admitted`. +It then snapshots nonresident waiting requests from the scheduler pool. These are `Assigned` and are marked as newly admitted candidates. The scheduler orders candidates with decodes first. Order remains stable among decodes and among prefills. Each candidate initially contributes one provisional @@ -617,11 +617,6 @@ Requests skipped because of token, row, or temporary cache capacity remain pendi If no request can run because of temporary capacity, `StepDynamic()` reports `CapacityDeferred` instead of returning `nullptr`. Returning `nullptr` would incorrectly tell the caller that no work remains. -The native Engine exposes `CapacityDeferred` as a structured `StepOutcomeKind`. The current C and -Python wrappers still surface it as an error message. Phase 2 must add a structured public -backpressure/residency signal before introducing automatic eviction, so applications can select a -turn-complete request to close without parsing text. - ## Static engine path The static engine path is intentionally separate. @@ -635,9 +630,7 @@ batch. Static cache rows still cannot be released independently, and an all-turn be recycled for new work. Static continuation is therefore valid only while the original single-request batch remains resident. -A closed static row remains physically retained until that shared batch is recycled. It is not -sampled or returned again, but its Request/Search storage can remain alive for the lifetime of the -batch. +A closed request that is already resident in a static batch remains physically retained until that shared batch is recycled. It is not sampled or returned again, but its Request/Search storage can remain alive for the lifetime of the batch. Changes to shared types such as `Request`, `ScheduledRequests`, `ModelExecutor`, or `SimpleDecoder` should be checked against both paths. This document should be updated only where behavior is shared or where the dynamic path changes. @@ -663,11 +656,6 @@ if request.status == og.RequestStatus.TURN_COMPLETE: engine.remove_request(request) ``` -`AddTokens` is for initial input. The explicit continuation operations are `OgaRequestContinue` in -C, `OgaRequest::Continue` in the C++ wrapper, and `request.continue_with` in Python. Lifecycle is -available through `OgaRequestGetStatus`, `OgaRequest::GetStatus`, and `request.status`. -`IsDone()` remains a compatibility convenience for β€œthe current turn is complete.” - One ready request may be returned several times over its lifetime as new tokens become available. A turn-complete dynamic request remains cache-resident until explicit removal, which releases dynamic cache ownership immediately. diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp index 367f54060a..b6c39f33e8 100644 --- a/src/engine/engine.cpp +++ b/src/engine/engine.cpp @@ -58,7 +58,6 @@ void Engine::AddRequest(std::shared_ptr request) { if (cache_manager_->SupportsDynamicBatching()) { request->ValidateEngineCompatibility(); } - scheduler_->ValidateRequest(*request); request->Assign(shared_from_this()); scheduler_->AddRequest(request); } @@ -73,16 +72,11 @@ void Engine::RemoveRequest(std::shared_ptr request) { scheduler_->RemoveRequest(request); - auto first_undrained = - ready_requests_.begin() + static_cast(ready_request_index_); - const auto retained_end = - std::remove(first_undrained, ready_requests_.end(), request); - const auto new_end = ready_request_index_ == 0 - ? retained_end - : std::move(first_undrained, retained_end, - ready_requests_.begin()); - ready_requests_.erase(new_end, ready_requests_.end()); + ready_requests_.erase( + ready_requests_.begin(), + ready_requests_.begin() + static_cast(ready_request_index_)); ready_request_index_ = 0; + std::erase(ready_requests_, request); request->CompleteClose(); } diff --git a/src/engine/request.cpp b/src/engine/request.cpp index d9be5d1dc3..60a2c03ecf 100644 --- a/src/engine/request.cpp +++ b/src/engine/request.cpp @@ -126,12 +126,12 @@ void Request::AddTokens(std::span tokens) { } void Request::Continue(std::span tokens) { + if (IsClosed(status_)) { + throw std::runtime_error("Cannot continue a closed request."); + } if (tokens.empty()) throw std::runtime_error("Expected at least one token for continuation. Received 0."); if (!IsTurnComplete(status_)) { - if (IsClosed(status_)) { - throw std::runtime_error("Cannot continue a closed request."); - } throw std::runtime_error("Continue is only valid after the current turn is complete."); } @@ -221,7 +221,7 @@ void Request::AdvanceChunk() { } int32_t Request::UnseenToken() { - if (next_unseen_token_index_ == unseen_token_indices_.size()) + if (next_unseen_token_index_ >= unseen_token_indices_.size()) throw std::runtime_error("All tokens have been seen."); const size_t token_index = unseen_token_indices_[next_unseen_token_index_++]; diff --git a/src/engine/scheduler.cpp b/src/engine/scheduler.cpp index 2315066b14..1ed3c99f8b 100644 --- a/src/engine/scheduler.cpp +++ b/src/engine/scheduler.cpp @@ -31,16 +31,13 @@ ScheduledRequests Scheduler::CreateScheduledRequests(const StepPlan& plan) { StaticBatchScheduler::StaticBatchScheduler(std::shared_ptr model, std::shared_ptr cache_manager) : Scheduler{model}, model_{model}, cache_manager_{cache_manager} {} -void StaticBatchScheduler::ValidateRequest(const Request& request) const { +void StaticBatchScheduler::AddRequest(std::shared_ptr request) { // The static batch decoder rebuilds its contiguous cache from the whole sequence every step, so it // cannot resume a half written prompt. Only the paged cache can hold one. - if (request.SearchOptions().chunk_size.value_or(0) != 0) { + if (request->SearchOptions().chunk_size.value_or(0) != 0) { throw std::runtime_error( "search.chunk_size requires dynamic batching; the static batch scheduler cannot chunk a prefill."); } -} - -void StaticBatchScheduler::AddRequest(std::shared_ptr request) { if (auto* sampler = GetBatchedSampler()) request->SamplingState(*sampler); requests_pool_.push_back(request); @@ -59,10 +56,6 @@ void StaticBatchScheduler::RemoveRequest(std::shared_ptr request) { ScheduledRequests StaticBatchScheduler::Schedule() { const auto allocated_requests = cache_manager_->AllocatedRequests(); - const auto is_resident = [&allocated_requests](const std::shared_ptr& request) { - return std::find(allocated_requests.begin(), allocated_requests.end(), request) != - allocated_requests.end(); - }; for (const auto& request : allocated_requests) { if (IsQueued(request->status_)) { @@ -72,7 +65,7 @@ ScheduledRequests StaticBatchScheduler::Schedule() { std::vector> requests_to_schedule; for (auto& request : requests_pool_) { - if (IsQueued(request->status_) && !is_resident(request)) { + if (IsQueued(request->status_) && !cache_manager_->IsResident(request)) { requests_to_schedule.push_back(request); } } @@ -196,10 +189,6 @@ StepPlanningResult DynamicBatchScheduler::PlanStep(StepPlan& plan) { candidates.push_back(std::move(candidate)); }; - const auto is_resident = [&allocated_requests](const std::shared_ptr& request) { - return std::find(allocated_requests.begin(), allocated_requests.end(), request) != - allocated_requests.end(); - }; for (const auto& request : allocated_requests) { if (IsTurnComplete(request->status_)) { continue; @@ -208,7 +197,7 @@ StepPlanningResult DynamicBatchScheduler::PlanStep(StepPlan& plan) { } for (const auto& request : requests_pool_) { - if (IsQueued(request->status_) && !is_resident(request)) { + if (IsQueued(request->status_) && !cache_manager_->IsResident(request)) { add_candidate(request, true); } } diff --git a/src/engine/scheduler.h b/src/engine/scheduler.h index 5e1080d03b..6d76fc9060 100644 --- a/src/engine/scheduler.h +++ b/src/engine/scheduler.h @@ -26,8 +26,6 @@ struct Scheduler { static std::unique_ptr Create(std::shared_ptr model, std::shared_ptr cache_manager); - virtual void ValidateRequest(const Request&) const {} - /** * @brief Adds a request to the Scheduler for processing. * @param request A shared pointer to the Request object to be added. @@ -84,8 +82,6 @@ struct Scheduler { struct StaticBatchScheduler : Scheduler { StaticBatchScheduler(std::shared_ptr model, std::shared_ptr cache_manager); - void ValidateRequest(const Request& request) const override; - void AddRequest(std::shared_ptr request) override; void RemoveRequest(std::shared_ptr request) override; diff --git a/test/engine/engine_invariants_tests.cpp b/test/engine/engine_invariants_tests.cpp index e99fb9fbd4..36f771d75d 100644 --- a/test/engine/engine_invariants_tests.cpp +++ b/test/engine/engine_invariants_tests.cpp @@ -309,14 +309,14 @@ TEST(InvariantValidatorTest, SeenBeyondCurrentReported) { EXPECT_FALSE(ValidateRequestInvariants(request).empty()); } -TEST(InvariantValidatorTest, CompletedRequestWithFinalUnprocessedTokenIsValid) { +TEST(InvariantValidatorTest, TurnCompleteRequestWithFinalUnprocessedTokenIsValid) { // At completion the just-generated final token is appended but never fed back to the model, so a - // A TurnComplete Request may legitimately report unprocessed tokens. This must not fire. + // TurnComplete Request may legitimately report unprocessed tokens. This must not fire. auto request = MakeValidRequest(kRequestA, RequestStatus::TurnComplete, 10, 9, 10); EXPECT_TRUE(ValidateRequestInvariants(request).empty()); } -TEST(InvariantValidatorTest, CompletedRequestFullyProcessedIsValid) { +TEST(InvariantValidatorTest, TurnCompleteRequestFullyProcessedIsValid) { EXPECT_TRUE(ValidateRequestInvariants( MakeValidRequest(kRequestA, RequestStatus::TurnComplete, 10, 10, 10)) .empty()); diff --git a/test/engine/engine_step_tests.cpp b/test/engine/engine_step_tests.cpp index 97465e3ece..8d277897c5 100644 --- a/test/engine/engine_step_tests.cpp +++ b/test/engine/engine_step_tests.cpp @@ -158,7 +158,7 @@ TEST_F(EngineStepTest, StepWithNoRequestsReturnsNull) { EXPECT_EQ(engine.executor->decode_calls, 0); } -TEST_F(EngineStepTest, StaticBatchingRetainsLegacyCommitOrdering) { +TEST_F(EngineStepTest, StaticBatchingPreservesOrderingAndReusesResidentContinuation) { model_->config_->engine.dynamic_batching.reset(); auto trace = std::make_shared(); auto cache = std::make_shared( @@ -190,28 +190,6 @@ TEST_F(EngineStepTest, StaticBatchingRetainsLegacyCommitOrdering) { EXPECT_EQ(cache_observer->allocate_calls, allocations_before); } -TEST_F(EngineStepTest, StaticAdmissionValidationDoesNotStrandRequest) { - model_->config_->engine.dynamic_batching.reset(); - auto cache = std::make_shared( - model_, /*capacity=*/4, nullptr, /*supports_dynamic_batching=*/false); - auto scheduler = Scheduler::Create(model_, cache); - auto executor = std::make_unique( - model_, cache, EosToken(*model_)); - EngineDependencies dependencies{cache, std::move(scheduler), - std::move(executor)}; - auto engine = std::make_shared(model_, std::move(dependencies)); - auto prompt = Prompt(10); - auto request = MintRequest(*model_, prompt); - request->Params()->search.chunk_size = 2; - - EXPECT_THROW(engine->AddRequest(request), std::runtime_error); - EXPECT_EQ(request->status_, RequestStatus::Unassigned); - - request->Params()->search.chunk_size.reset(); - EXPECT_NO_THROW(engine->AddRequest(request)); - EXPECT_EQ(request->status_, RequestStatus::Assigned); -} - TEST_F(EngineStepTest, StaticContinueFailsAfterBatchRecycling) { model_->config_->engine.dynamic_batching.reset(); auto cache = std::make_shared(model_); @@ -529,27 +507,6 @@ TEST_F(EngineStepTest, RemovingOnlyDrainedReadyRequestClearsQueue) { EXPECT_EQ(engine.engine->Step(), nullptr); } -TEST_F(EngineStepTest, RepeatedReadyRemovalsPreserveRemainingQueueOrder) { - auto engine = MakeDoublesEngine(model_, /*capacity=*/8, /*forced_token=*/5); - std::vector> requests; - for (int32_t seed : {10, 20, 30}) { - auto prompt = Prompt(seed); - auto request = MintRequest(*model_, prompt); - engine.engine->AddRequest(request); - requests.push_back(std::move(request)); - } - - ASSERT_EQ(engine.engine->Step(), requests[0]); - engine.engine->RemoveRequest(requests[0]); - // The first removal resets the drain cursor to zero while two ready entries remain. Removing a - // second request must compact that nonempty queue in place without overlapping std::move ranges. - engine.engine->RemoveRequest(requests[1]); - - EXPECT_EQ(requests[0]->status_, RequestStatus::Closed); - EXPECT_EQ(requests[1]->status_, RequestStatus::Closed); - EXPECT_EQ(engine.engine->Step(), requests[2]); -} - TEST_F(EngineStepTest, StaticTurnCompleteRowIsNotRepublishedWhilePeerRuns) { model_->config_->engine.dynamic_batching.reset(); auto cache = std::make_shared( diff --git a/test/engine/request_lifecycle_tests.cpp b/test/engine/request_lifecycle_tests.cpp index 2e2b38bf61..6680aaabb9 100644 --- a/test/engine/request_lifecycle_tests.cpp +++ b/test/engine/request_lifecycle_tests.cpp @@ -130,14 +130,12 @@ TEST_F(RequestLifecycleTest, AppendIsRejectedWhileInProgress) { EXPECT_EQ(request->CurrentSequenceLength(), length_before); } -TEST_F(RequestLifecycleTest, AddTokensIsRejectedAfterSubmission) { +TEST_F(RequestLifecycleTest, AddTokensIsRejectedWhileAssigned) { auto prompt = Prompt(); const std::vector more{5}; auto request = MintAssignedRequest(engine_.engine, *model_, prompt); EXPECT_THROW(request->AddTokens(more), std::runtime_error); - request->Schedule(); - EXPECT_THROW(request->AddTokens(more), std::runtime_error); } // After a turn completes, Continue appends another input fragment and queues the resident request. diff --git a/test/engine/scheduler_contract_tests.cpp b/test/engine/scheduler_contract_tests.cpp index 66b4536c3e..3f522e33d9 100644 --- a/test/engine/scheduler_contract_tests.cpp +++ b/test/engine/scheduler_contract_tests.cpp @@ -221,26 +221,6 @@ TEST_F(SchedulerContractTest, DynamicResidentQueuedRequestIsNotReadmitted) { EXPECT_EQ(cache->AllocatedCount(), 1u); } -TEST_F(SchedulerContractTest, StaticResidentQueuedRequestSchedulesWithoutAllocation) { - model_->config_->engine.dynamic_batching.reset(); - auto cache = std::make_shared( - model_, /*capacity=*/4, nullptr, /*supports_dynamic_batching=*/false); - StaticBatchScheduler scheduler(model_, cache); - - auto request = Assigned(10); - scheduler.AddRequest(request); - cache->Allocate({request}); - request->status_ = RequestStatus::Assigned; - const int allocations_before = cache->allocate_calls; - - auto scheduled = scheduler.Schedule(); - - ASSERT_EQ(scheduled.size(), 1u); - EXPECT_EQ(scheduled[0], request); - EXPECT_EQ(request->status_, RequestStatus::InProgress); - EXPECT_EQ(cache->allocate_calls, allocations_before); -} - // Removing a request from the dynamic scheduler deallocates its cache resources immediately. TEST_F(SchedulerContractTest, DynamicRemoveReleasesCacheResources) { auto cache = std::make_shared(model_, /*capacity=*/8); diff --git a/test/python/test_onnxruntime_genai_engine.py b/test/python/test_onnxruntime_genai_engine.py index 999d83b0e5..4117d01652 100644 --- a/test/python/test_onnxruntime_genai_engine.py +++ b/test/python/test_onnxruntime_genai_engine.py @@ -250,14 +250,6 @@ def test_continuation_while_peer_remains_active(model): assert short_sink.tokens == reference_sink.tokens -def test_request_rejects_empty_input(model): - params = og.GeneratorParams(model) - request = og.Request(params) - - with pytest.raises(RuntimeError, match="at least one token"): - request.add_tokens(np.asarray([], dtype=np.int32)) - - def test_request_rejects_tokens_while_awaiting_admission(model): engine = og.Engine(model) sink = _Sink() @@ -269,19 +261,6 @@ def test_request_rejects_tokens_while_awaiting_admission(model): engine.remove_request(request) -def test_closed_request_cannot_continue(model): - engine = og.Engine(model) - sink = _Sink() - request = _add_request(engine, model, _PROMPT_A, 8, sink) - engine.remove_request(request) - - assert request.status == og.RequestStatus.CLOSED - with pytest.raises(RuntimeError, match="closed request"): - request.add_tokens(np.asarray([12], dtype=np.int32)) - with pytest.raises(RuntimeError, match="closed request"): - request.continue_with(np.asarray([12], dtype=np.int32)) - - def test_request_cannot_be_removed_from_another_engine(model): owner = og.Engine(model) other = og.Engine(model) @@ -320,6 +299,10 @@ def test_request_lifecycle_status(model): engine.remove_request(request) assert request.status == og.RequestStatus.CLOSED + with pytest.raises(RuntimeError, match="closed request"): + request.add_tokens(np.asarray([12], dtype=np.int32)) + with pytest.raises(RuntimeError, match="closed request"): + request.continue_with(np.asarray([12], dtype=np.int32)) with pytest.raises(RuntimeError, match="already closed"): engine.remove_request(request) From 0467ed40ad5976a2cc481b120fa733612e24a065 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 15:39:29 -0500 Subject: [PATCH 03/19] Clarify retained static-row sampling behavior Use TurnComplete, Closed, and InProgress terminology so the decoder comment matches the lifecycle predicates used by ScheduledRequests. Verified at src/engine/decoders/static_batch_decoder_io.cpp:150-158. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b232f82c-f25c-429a-a855-7f8c8f40bbf3 --- src/engine/decoders/static_batch_decoder_io.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/engine/decoders/static_batch_decoder_io.cpp b/src/engine/decoders/static_batch_decoder_io.cpp index 4279660ece..e6288ff4c2 100644 --- a/src/engine/decoders/static_batch_decoder_io.cpp +++ b/src/engine/decoders/static_batch_decoder_io.cpp @@ -150,9 +150,9 @@ void StaticBatchDecoderIO::PrepareLogits(std::shared_ptr mode std::vector> StaticBatchDecoderIO::ProcessLogits() { std::vector valid_token_indices; for (auto& request : scheduled_requests_) { - // A completed row retained in the static batch for continuation can contribute zero + // A TurnComplete or Closed row retained in the static batch can contribute zero // unprocessed tokens. Selecting index 0 keeps the subspan in bounds; its logits are discarded - // because ScheduledRequests skips completed and removed rows during sampling. + // because ScheduledRequests samples only InProgress rows. const auto unprocessed_token_count = request->ScheduledTokenCount(); valid_token_indices.push_back( unprocessed_token_count == 0 ? 0 : static_cast(unprocessed_token_count - 1)); From e8e5d8f2dfd05333a1b1111df4b64fb183b40d38 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 18:09:58 -0500 Subject: [PATCH 04/19] Preserve DML chat with token-history replay Recreate DML Engine requests from the exact accumulated logical token history while continuation-capable providers retain resident request reuse. File: examples/python/engine/model-qa.py Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b232f82c-f25c-429a-a855-7f8c8f40bbf3 --- examples/python/engine/model-qa.py | 39 ++++++++++++++++++++++++------ 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/examples/python/engine/model-qa.py b/examples/python/engine/model-qa.py index 1e02ffd19f..6ceaa61b72 100644 --- a/examples/python/engine/model-qa.py +++ b/examples/python/engine/model-qa.py @@ -23,13 +23,23 @@ def run(args: argparse.Namespace): max_length=1024, ) - request = og.Request(params) system_message = json.dumps([{"role": "system", "content": ""}]) - request.add_tokens( - tokenizer.encode( - tokenizer.apply_chat_template(messages=system_message, add_generation_prompt=False), - ), + system_tokens = tokenizer.encode( + tokenizer.apply_chat_template(messages=system_message, add_generation_prompt=False), ) + # Temporary DML-only fallback while the low-level Engine continuation API is transitional: + # replay exact token IDs in a fresh request instead of reconstructing text. Replaying the full + # session also preserves max_length as a session-total limit. Other advertised providers keep + # one resident request and use continue_with(). + use_dml_replay = args.execution_provider == "dml" + logical_token_history = [int(token) for token in system_tokens] + eos_token_ids = {int(token) for token in tokenizer.eos_token_ids} + + request = None + if not use_dml_replay: + request = og.Request(params) + request.add_tokens(system_tokens) + streaming_tokenizer = tokenizer.create_stream() request_added = False @@ -43,7 +53,13 @@ def run(args: argparse.Namespace): tokenizer.apply_chat_template(messages=user_message, add_generation_prompt=True), ) - if request_added: + if use_dml_replay: + logical_token_history.extend(int(token) for token in turn_tokens) + request = og.Request(params) + request.add_tokens(logical_token_history) + engine.add_request(request) + request_added = True + elif request_added: request.continue_with(turn_tokens) else: request.add_tokens(turn_tokens) @@ -54,15 +70,22 @@ def run(args: argparse.Namespace): while ready_request := engine.step(): while ready_request.has_unseen_tokens(): + token = ready_request.get_unseen_token() + if use_dml_replay and token not in eos_token_ids: + logical_token_history.append(token) print( - streaming_tokenizer.decode(ready_request.get_unseen_token()), + streaming_tokenizer.decode(token), end="", flush=True, ) print() + if use_dml_replay: + engine.remove_request(request) + request_added = False + request = None finally: - if request_added: + if request_added and request is not None: engine.remove_request(request) From b728ca0a12e6a4ef560a9c2d123b5a4ed84d34a0 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 18:56:32 -0500 Subject: [PATCH 05/19] Finalize active request lifecycle and orphan cleanup Rename the internal and public executing state to Active, add precise turn-completion APIs, and deprecate ambiguous Request IsDone aliases. Make terminal removal idempotent, reject continuation with an undrained ready event, and defer automatic reclamation of requests whose final external handle is released. Document the transitional serialized API, cumulative max_length, ordered untagged output, and immediate-versus-deferred cleanup contracts. Files: src/engine/**; src/smartptrs.h; src/ort_genai*.{h,cpp}; src/python/python.cpp; docs/paged_attention_engine.md; focused C/C++/Python lifecycle tests. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b232f82c-f25c-429a-a855-7f8c8f40bbf3 --- docs/paged_attention_engine.md | 48 ++-- .../decoders/static_batch_decoder_io.cpp | 2 +- src/engine/engine.cpp | 57 +++- src/engine/engine.h | 2 + src/engine/request.cpp | 28 +- src/engine/request.h | 19 +- src/engine/request_status.h | 4 +- src/engine/scheduled_requests.cpp | 2 +- src/ort_genai.h | 29 ++- src/ort_genai_c.cpp | 12 +- src/ort_genai_c.h | 53 +++- src/python/python.cpp | 20 +- src/smartptrs.h | 22 +- test/c_api_tests.cpp | 27 ++ test/engine/engine_invariants_tests.cpp | 16 +- test/engine/engine_step_tests.cpp | 243 +++++++++++++++++- test/engine/request_lifecycle_tests.cpp | 57 +++- test/engine/scheduler_contract_tests.cpp | 2 +- test/python/test_onnxruntime_genai_engine.py | 42 ++- 19 files changed, 586 insertions(+), 99 deletions(-) diff --git a/docs/paged_attention_engine.md b/docs/paged_attention_engine.md index 81b40fd3d6..2da0682f64 100644 --- a/docs/paged_attention_engine.md +++ b/docs/paged_attention_engine.md @@ -10,6 +10,10 @@ The `Engine` can use either static batching or dynamic batching. This document f The current dynamic path manages paged KV decoder state together with per-request search and sampler state. It does not bind or transactionally checkpoint hybrid recurrent, convolutional, or other mutable model state. Future support for such state must be selected from model capabilities rather than model names, and every Engine-owned mutable state must participate in the same transaction boundary. +> **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 status 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. + The main implementation is under `src/engine/`: | Responsibility | Main files | @@ -92,13 +96,13 @@ The engine creates throughput by batching several independent requests, not by p The important request states are: ```text -Unassigned (Created) -- submit --> Assigned (Queued) -- schedule --> InProgress +Unassigned (Created) -- submit --> Assigned (Queued) -- schedule --> Active ^ | | | turn stops +---- Continue(tokens) ---- TurnComplete Assigned (Queued) ---+ -InProgress ----------+-- Remove() --> Closed +Active --------------+-- Remove() --> Closed TurnComplete --------+ ``` @@ -120,7 +124,9 @@ initializes the sequence counters, and records the owning Engine. `AddTokens()` both rejected while already queued. Input must leave room for at least one generated token below `max_length`. -### `InProgress` +`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. + +### `Active` The current turn is executable and owned by the Engine. It normally has one unprocessed token at the beginning of a decode step: the token sampled by the previous step. @@ -128,7 +134,7 @@ the beginning of a decode step: the token sampled by the previous step. ### `TurnComplete` The current generation turn reached an end condition, such as EOS or maximum length. Generated -output remains available, and `IsDone()` means this state rather than permanent request termination. +output remains available. `IsTurnComplete()` is the precise API for testing this state. `IsDone()` is a temporary compatibility alias for `IsTurnComplete()`; it remains turn-scoped and does not mean permanent request termination or removal. A generated EOS/stop token is not appended to the logical sequence or returned as unseen output. The next continuation fragment is therefore responsible for any turn-boundary tokens required by @@ -137,17 +143,21 @@ the model's chat template. `Continue(tokens)` appends the next input fragment and moves a resident request back to `Assigned`. `AddTokens()` remains an initial-input-only operation. +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. + 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. +longer need continuation when deterministic immediate reclamation is required. ### `Remove()` -`Remove()` is legal from `Assigned`, `InProgress`, and `TurnComplete`, and moves the request to +`Remove()` is legal from `Assigned`, `Active`, and `TurnComplete`, and moves the request to terminal `Closed`. On the dynamic path, removal immediately erases scheduler membership and releases committed paged-cache ownership. The Engine also removes any undrained ready-queue entries for that request. +Calling `Remove()` for an already terminal `Closed` request is an engine-agnostic idempotent no-op because the request no longer has an owner. Removing an `Unassigned` request remains invalid, as does asking an Engine other than the request's owner to remove a nonterminal request. This idempotence does not relax the external serialization requirement. + ### `Closed` `Closed` is distinct from `Unassigned` because removal may already have destroyed residency and @@ -174,6 +184,8 @@ logical sequence: | `unseen_token_indices_` | Positions of generated tokens in `tokens_host_`; continuation-input positions are never added | | `next_unseen_token_index_` | Cursor into `unseen_token_indices_`; entries at and after this cursor have not been consumed | +Unread generated output is one globally ordered stream for the request across all turns. The unseen-output API does not tag tokens with a turn, so callers that need per-turn attribution must track the boundaries themselves. + The unprocessed tokens are: ```text @@ -208,7 +220,7 @@ This separation between search length and processed length is what lets each mod ## `Engine::Step()` and ready-result draining -The public engine API advances through repeated calls to `Step()`. +The current transitional low-level engine API advances through repeated calls to `Step()`. Before executing new work, `Step()` checks `ready_requests_`. One model invocation may produce a token for several requests, but `Step()` returns only one `Request` pointer. The remaining ready requests stay in the ready queue and are returned by later `Step()` calls without another model execution. @@ -237,7 +249,7 @@ application decides which conversation to release with `Remove()`. ### 2. Build the initial step plan The scheduler snapshots requests that already belong to the paged cache. Executable residents may -be `InProgress` or `Assigned`; an `Assigned` resident is a queued continuation. +be `Active` or `Assigned`; an `Assigned` resident is a queued continuation. It then snapshots nonresident waiting requests from the scheduler pool. These are `Assigned` and are marked as newly admitted candidates. @@ -281,7 +293,7 @@ admissions remains within `max_batch_size`. Selected entries are compacted to the beginning of the plan. Deferred entries remain unchanged in committed state: -- An existing request keeps its current block table and stays `InProgress`. +- An existing request keeps its current block table and stays `Active`. - A new request stays `Assigned` and owns no cache blocks. - Both can be considered again by the next engine step. @@ -431,10 +443,10 @@ Committing request bookkeeping: - Appends the staged token to the host token mirror. - Sets `processed_sequence_length_` to the sequence length that existed before sampling. -- Changes the status to `InProgress` or `TurnComplete`. +- Changes the status to `Active` or `TurnComplete`. For a new request or queued continuation, this commit is the point where it moves from `Assigned` -to `InProgress` or `TurnComplete`. The dynamic transaction path does not need a separate visible +to `Active` or `TurnComplete`. The dynamic transaction path does not need a separate visible scheduling state between those states. Finally, the engine swaps the staged ready list into `ready_requests_`. The first ready request is returned immediately, and later calls drain the rest without another model run. @@ -560,7 +572,7 @@ from taking the capacity needed to finish an admitted partial prefill. ### Decode -An in-progress request normally contributes one token: the token sampled by its previous committed step. +An active request normally contributes one token: the token sampled by its previous committed step. ### Mixed batch @@ -625,7 +637,7 @@ The static engine path is intentionally separate. `StepStatic()` performs decode and sampling directly without the dynamic transaction and reservation protocol. -A resident static request queued by `Continue()` returns to `InProgress` without reallocating the +A resident static request queued by `Continue()` returns to `Active` without reallocating the batch. Static cache rows still cannot be released independently, and an all-turn-complete batch may be recycled for new work. Static continuation is therefore valid only while the original single-request batch remains resident. @@ -634,9 +646,9 @@ A closed request that is already resident in a static batch remains physically r Changes to shared types such as `Request`, `ScheduledRequests`, `ModelExecutor`, or `SimpleDecoder` should be checked against both paths. This document should be updated only where behavior is shared or where the dynamic path changes. -## Public API shape +## Transitional low-level public API shape -The language bindings expose the same basic loop: +The language bindings currently expose the same basic low-level loop. Production hosts are expected to wrap this surface or use its replacement rather than expose it as their stable API: ```python request.add_tokens(initial_tokens) @@ -649,16 +661,18 @@ while engine.has_pending_requests(): token = ready_request.get_unseen_token() # Stream or process the token. -if request.status == og.RequestStatus.TURN_COMPLETE: +if request.is_turn_complete(): request.continue_with(next_turn_tokens) # Repeat engine.step(), then close the conversation when continuation is no longer needed. +# Explicit removal releases resources immediately; final-handle release otherwise defers cleanup +# until the next add_request() or step() boundary. engine.remove_request(request) ``` One ready request may be returned several times over its lifetime as new tokens become available. A turn-complete dynamic request remains cache-resident until explicit removal, which releases dynamic -cache ownership immediately. +cache ownership immediately. The unseen-output accessors return generated tokens one at a time in global request order without turn tags. ## Keeping this document current diff --git a/src/engine/decoders/static_batch_decoder_io.cpp b/src/engine/decoders/static_batch_decoder_io.cpp index e6288ff4c2..6ffa0ea024 100644 --- a/src/engine/decoders/static_batch_decoder_io.cpp +++ b/src/engine/decoders/static_batch_decoder_io.cpp @@ -152,7 +152,7 @@ std::vector> StaticBatchDecoderIO::ProcessLogits() { for (auto& request : scheduled_requests_) { // A TurnComplete or Closed row retained in the static batch can contribute zero // unprocessed tokens. Selecting index 0 keeps the subspan in bounds; its logits are discarded - // because ScheduledRequests samples only InProgress rows. + // because ScheduledRequests samples only Active rows. const auto unprocessed_token_count = request->ScheduledTokenCount(); valid_token_indices.push_back( unprocessed_token_count == 0 ? 0 : static_cast(unprocessed_token_count - 1)); diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp index b6c39f33e8..1e56b66ecc 100644 --- a/src/engine/engine.cpp +++ b/src/engine/engine.cpp @@ -55,16 +55,27 @@ EngineDependencies Engine::CreateDependencies(std::shared_ptr model) { } void Engine::AddRequest(std::shared_ptr request) { + ReclaimAbandonedRequests(); if (cache_manager_->SupportsDynamicBatching()) { request->ValidateEngineCompatibility(); } - request->Assign(shared_from_this()); - scheduler_->AddRequest(request); + + // 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; + } } void Engine::RemoveRequest(std::shared_ptr request) { if (request && IsClosed(request->status_)) { - throw std::runtime_error("Cannot remove a request that is already closed."); + return; } if (!request || request->engine_.lock().get() != this) { throw std::runtime_error("Cannot remove a request from an engine it does not belong to."); @@ -77,7 +88,39 @@ void Engine::RemoveRequest(std::shared_ptr request) { ready_requests_.begin() + static_cast(ready_request_index_)); ready_request_index_ = 0; std::erase(ready_requests_, request); + std::erase(staged_ready_requests_, request); request->CompleteClose(); + std::erase_if(tracked_requests_, [&request](const std::weak_ptr& tracked) { + const auto owned = tracked.lock(); + return !owned || owned == 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 + // release, ready-notification purge, and terminal close. + std::vector> abandoned_requests; + abandoned_requests.reserve(tracked_requests_.size()); + std::erase_if(tracked_requests_, [&abandoned_requests, this](const std::weak_ptr& tracked) { + const auto request = tracked.lock(); + if (!request) { + return true; + } + if (!IsClosed(request->status_) && + request->engine_.lock().get() == this && + request->IsExternallyAbandoned()) { + abandoned_requests.push_back(request); + } + return false; + }); + + for (const auto& request : abandoned_requests) { + // Recheck defensively in case an external owner was reacquired before this serialized boundary. + if (request->IsExternallyAbandoned()) { + RemoveRequest(request); + } + } } void Engine::ValidateRequestCanContinue(const std::shared_ptr& request) const { @@ -92,6 +135,13 @@ void Engine::ValidateRequestCanContinue(const std::shared_ptr& request) throw std::runtime_error("Cannot continue a request whose model state is no longer resident."); } + if (std::find(ready_requests_.begin() + static_cast(ready_request_index_), + ready_requests_.end(), request) != ready_requests_.end()) { + throw std::runtime_error( + "Cannot continue a request while its ready notification is pending; " + "call Engine::Step() to drain the ready notification before continuing."); + } + if (!cache_manager_->SupportsDynamicBatching() && cache_manager_->ResidentRequestCount() > 1) { throw std::runtime_error( @@ -100,6 +150,7 @@ void Engine::ValidateRequestCanContinue(const std::shared_ptr& request) } std::shared_ptr Engine::Step() { + ReclaimAbandonedRequests(); if (auto request = DrainReadyRequest()) { return request; } diff --git a/src/engine/engine.h b/src/engine/engine.h index 0e734e7f3d..a8bca9718a 100644 --- a/src/engine/engine.h +++ b/src/engine/engine.h @@ -119,6 +119,7 @@ struct Engine : std::enable_shared_from_this, bool HasPendingRequests() const; private: + void ReclaimAbandonedRequests(); std::shared_ptr DrainReadyRequest(); std::shared_ptr StepDynamic(); std::shared_ptr StepStatic(); @@ -139,6 +140,7 @@ struct Engine : std::enable_shared_from_this, EngineTransactionMetrics transaction_metrics_; StepPlan step_plan_; std::vector step_results_; + std::vector> tracked_requests_; std::vector> ready_requests_; std::vector> staged_ready_requests_; size_t ready_request_index_{}; diff --git a/src/engine/request.cpp b/src/engine/request.cpp index 60a2c03ecf..c352b3818a 100644 --- a/src/engine/request.cpp +++ b/src/engine/request.cpp @@ -53,6 +53,18 @@ Request::Request(std::shared_ptr params) search_->DeferCompletion(true); } +void Request::OnFirstExternalReference() noexcept { + externally_abandoned_.store(false, std::memory_order_release); +} + +void Request::OnLastExternalReference() noexcept { + externally_abandoned_.store(true, std::memory_order_release); +} + +bool Request::IsExternallyAbandoned() const noexcept { + return externally_abandoned_.load(std::memory_order_acquire); +} + void Request::Assign(std::shared_ptr engine) { if (status_ != RequestStatus::Unassigned) { throw std::runtime_error("Cannot add the request to the engine since it is already assigned."); @@ -83,7 +95,7 @@ void Request::Schedule() { throw std::runtime_error("Cannot schedule a request with no tokens."); } - status_ = RequestStatus::InProgress; + status_ = RequestStatus::Active; } void Request::Remove() { @@ -91,7 +103,7 @@ void Request::Remove() { throw std::runtime_error("Cannot close a request that has not been submitted to an engine."); } if (IsClosed(status_)) { - throw std::runtime_error("Cannot close a request that is already closed."); + return; } auto engine = engine_.lock(); @@ -112,7 +124,7 @@ void Request::AddTokens(std::span tokens) { throw std::runtime_error("Expected at least one token for generation. Received 0."); if (status_ != RequestStatus::Unassigned) { - if (IsTurnComplete(status_)) { + if (IsTurnComplete()) { throw std::runtime_error("AddTokens only accepts initial input; use Continue for another turn."); } if (IsClosed(status_)) { @@ -131,7 +143,7 @@ void Request::Continue(std::span tokens) { } if (tokens.empty()) throw std::runtime_error("Expected at least one token for continuation. Received 0."); - if (!IsTurnComplete(status_)) { + if (!IsTurnComplete()) { throw std::runtime_error("Continue is only valid after the current turn is complete."); } @@ -254,10 +266,14 @@ std::span Request::UnprocessedTokensCpu() const { return std::span{tokens_host_}.subspan(begin, end - begin); } -bool Request::IsDone() const { +bool Request::IsTurnComplete() const { return status_ == RequestStatus::TurnComplete; } +bool Request::IsDone() const { + return IsTurnComplete(); +} + bool Request::IsPrefill() const { return processed_sequence_length_ < prompt_sequence_length_; } @@ -351,7 +367,7 @@ void Request::CommitStep(const RequestStepPlan& plan, unseen_token_indices_.push_back(token_index); } processed_sequence_length_ = static_cast(plan.target_cache_slots); - status_ = result.done ? RequestStatus::TurnComplete : RequestStatus::InProgress; + status_ = result.done ? RequestStatus::TurnComplete : RequestStatus::Active; } void Request::ApplyLogitsProcessors(DeviceSpan logits) { diff --git a/src/engine/request.h b/src/engine/request.h index fdf933f54a..6ebe15565d 100644 --- a/src/engine/request.h +++ b/src/engine/request.h @@ -16,6 +16,13 @@ namespace Generators { +struct Request; + +template <> +struct ExternalRefCountedTraits { + static constexpr bool notify_external_reference_changes = true; +}; + struct RequestStepResult { int32_t token{}; bool token_appended{}; @@ -50,7 +57,7 @@ struct Request : std::enable_shared_from_this, void Assign(std::shared_ptr engine); /** - * @brief Updates the status of the request to InProgress and prepares it for processing. + * @brief Updates the status of the request to Active and prepares it for processing. */ void Schedule(); @@ -149,6 +156,11 @@ struct Request : std::enable_shared_from_this, * @brief Checks if the current generation turn reached a stopping condition. * @return True in TurnComplete; the request may still be continued or closed. */ + bool IsTurnComplete() const; + + /** + * @brief Compatibility alias for IsTurnComplete(). + */ bool IsDone() const; RequestStatus Status() const noexcept { return status_; } @@ -297,8 +309,12 @@ struct Request : std::enable_shared_from_this, size_t next_unseen_token_index_{}; int64_t seen_sequence_length_{}; friend struct Engine; + friend struct ExternalRefCounted; void CompleteClose(); + void OnFirstExternalReference() noexcept; + void OnLastExternalReference() noexcept; + bool IsExternallyAbandoned() const noexcept; int64_t processed_sequence_length_{}; // Sequence length the application's tokens reach up to. Everything below it is prompt, so the @@ -309,6 +325,7 @@ struct Request : std::enable_shared_from_this, std::unique_ptr search_; std::unique_ptr batched_sampler_state_; std::weak_ptr engine_; + std::atomic externally_abandoned_{false}; void ApplyLogitsProcessors(DeviceSpan logits); void SelectNextToken(); diff --git a/src/engine/request_status.h b/src/engine/request_status.h index 8809c30173..766ab4080f 100644 --- a/src/engine/request_status.h +++ b/src/engine/request_status.h @@ -17,7 +17,7 @@ namespace Generators { enum class RequestStatus { Unassigned, // Created: initial input may be added before submission to an Engine. Assigned, // Queued: submitted initial work or a resident continuation awaits execution. - InProgress, // The current generation turn is executable and owned by the Engine. + Active, // The current generation turn is executable and owned by the Engine. TurnComplete, // The current turn stopped; output and resident model state remain available. Closed, // Permanently terminal; no scheduler or cache resources remain owned. }; @@ -27,7 +27,7 @@ constexpr bool IsQueued(RequestStatus status) noexcept { } constexpr bool IsExecuting(RequestStatus status) noexcept { - return status == RequestStatus::InProgress; + return status == RequestStatus::Active; } constexpr bool IsExecutable(RequestStatus status) noexcept { diff --git a/src/engine/scheduled_requests.cpp b/src/engine/scheduled_requests.cpp index c3cb824b0e..ad2041efa6 100644 --- a/src/engine/scheduled_requests.cpp +++ b/src/engine/scheduled_requests.cpp @@ -202,7 +202,7 @@ bool ScheduledRequests::PrepareBatchedSamplingPlan( sampling_plan_->Clear(); for (const auto& request : requests_) { // Dynamic transactions keep newly admitted and continued requests Queued until commit, while - // the static scheduler moves every executable row to InProgress before constructing the batch. + // the static scheduler moves every executable row to Active before constructing the batch. const bool status_is_executable = require_transaction_support ? IsExecutable(request->status_) : IsExecuting(request->status_); diff --git a/src/ort_genai.h b/src/ort_genai.h index 1e761b8251..04127834e0 100644 --- a/src/ort_genai.h +++ b/src/ort_genai.h @@ -903,10 +903,23 @@ struct OgaRequest : OgaAbstract { OgaCheckResult(OgaRequestContinue(this, &tokens)); } + /** + * \brief Returns whether the current generation turn is complete. + */ + bool IsTurnComplete() const { + bool is_turn_complete{}; + OgaCheckResult(OgaRequestIsTurnComplete(this, &is_turn_complete)); + return is_turn_complete; + } + + /** + * \brief Deprecated compatibility alias for IsTurnComplete(). + * + * \deprecated Use IsTurnComplete() instead. + */ + [[deprecated("Use IsTurnComplete() instead.")]] bool IsDone() const { - bool is_done{}; - OgaCheckResult(OgaRequestIsDone(this, &is_done)); - return is_done; + return IsTurnComplete(); } OgaRequestStatus GetStatus() const { @@ -953,10 +966,20 @@ struct OgaEngine : OgaAbstract { return f; } + /** + * \brief Submits a request and gives the engine ownership until Remove() is called. + * + * Ownership continues after the current turn completes. Remove the request before releasing its final handle. + */ void Add(OgaRequest& request) { OgaCheckResult(OgaEngineAddRequest(this, &request)); } + /** + * \brief Removes a request and releases engine ownership. + * + * Repeated calls after the request reaches OgaRequestStatus_closed are successful no-ops. + */ void Remove(OgaRequest& request) { OgaCheckResult(OgaEngineRemoveRequest(this, &request)); } diff --git a/src/ort_genai_c.cpp b/src/ort_genai_c.cpp index 18f5a7b95b..01e41ecd07 100644 --- a/src/ort_genai_c.cpp +++ b/src/ort_genai_c.cpp @@ -1361,13 +1361,17 @@ OgaResult* OgaRequestGetUnseenToken(OgaRequest* request, int32_t* token) { OGA_CATCH } -OgaResult* OgaRequestIsDone(const OgaRequest* request, bool* out) { +OgaResult* OgaRequestIsTurnComplete(const OgaRequest* request, bool* out) { OGA_TRY - *out = request->IsDone(); + *out = request->IsTurnComplete(); return nullptr; OGA_CATCH } +OgaResult* OgaRequestIsDone(const OgaRequest* request, bool* out) { + return OgaRequestIsTurnComplete(request, out); +} + OgaResult* OgaRequestGetStatus(const OgaRequest* request, OgaRequestStatus* out) { OGA_TRY switch (request->Status()) { @@ -1377,8 +1381,8 @@ OgaResult* OgaRequestGetStatus(const OgaRequest* request, OgaRequestStatus* out) case Generators::RequestStatus::Assigned: *out = OgaRequestStatus_queued; break; - case Generators::RequestStatus::InProgress: - *out = OgaRequestStatus_in_progress; + case Generators::RequestStatus::Active: + *out = OgaRequestStatus_active; break; case Generators::RequestStatus::TurnComplete: *out = OgaRequestStatus_turn_complete; diff --git a/src/ort_genai_c.h b/src/ort_genai_c.h index 6bf9ac40f0..a7c87dd909 100644 --- a/src/ort_genai_c.h +++ b/src/ort_genai_c.h @@ -62,7 +62,7 @@ typedef enum OgaElementType { typedef enum OgaRequestStatus { OgaRequestStatus_created, OgaRequestStatus_queued, - OgaRequestStatus_in_progress, + OgaRequestStatus_active, OgaRequestStatus_turn_complete, OgaRequestStatus_closed, } OgaRequestStatus; @@ -1164,6 +1164,8 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaCreateEngine(OgaModel* model, OgaEngine** /** * \brief Destroys the given engine. + * + * Remove every submitted request with OgaEngineRemoveRequest before destroying the engine. * \param[in] engine The engine to be destroyed. */ OGA_EXPORT void OGA_API_CALL OgaDestroyEngine(OgaEngine* engine); @@ -1181,7 +1183,9 @@ OGA_EXPORT void OGA_API_CALL OgaDestroyEngine(OgaEngine* engine); * * \param[in] engine The engine instance to run a processing step on. * \param[out] request A request that has been processed by the engine and is ready to be queried for results. - * If the engine has no ready requests, this will be set to a nullptr. + * If the engine has no ready requests, this will be set to a nullptr. Each non-null handle + * returned by this function must be released with OgaDestroyRequest. Releasing this handle + * does not remove the request from the engine. * \return OgaResult containing the error message if the operation failed, or nullptr on success. */ OGA_EXPORT OgaResult* OGA_API_CALL OgaEngineStep(OgaEngine* engine, OgaRequest** request); @@ -1190,6 +1194,8 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaEngineStep(OgaEngine* engine, OgaRequest** * \brief Checks if the engine has any pending requests to process. * * This function queries the OgaEngine to determine whether there are any requests that have not yet been fully processed. + * A false result does not mean the engine owns no requests: requests at OgaRequestStatus_turn_complete remain owned + * until explicitly removed with OgaEngineRemoveRequest. * * \param[in] engine The engine instance to check for pending requests. * \param[out] out Pointer to a boolean value that will be set to true if there are pending requests, or false otherwise. @@ -1202,9 +1208,13 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaEngineHasPendingRequests(OgaEngine* engine * * This function submits a new request to the engine, which will then be processed in subsequent calls to OgaEngineStep. * The request must be created using OgaCreateRequest and should contain the necessary parameters for model inference. + * On success, the engine retains ownership of the request at OgaRequestStatus_turn_complete, which completes only + * the current generation turn. OgaEngineRemoveRequest releases that ownership immediately. If the caller instead + * releases every external handle, the request is marked abandoned and reclaimed before the engine's next + * OgaEngineAddRequest or OgaEngineStep boundary. * * \param[in] engine The engine instance to which the request is being added. - * \param[in] request The request to add to the engine. The request must remain valid until it is removed or processed. + * \param[in] request The request to add to the engine. * \return OgaResult containing the error message if the operation failed, or nullptr on success. */ OGA_EXPORT OgaResult* OGA_API_CALL OgaEngineAddRequest(OgaEngine* engine, OgaRequest* request); @@ -1212,12 +1222,14 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaEngineAddRequest(OgaEngine* engine, OgaReq /** * \brief Removes a request from the OgaEngine. * - * This function removes a request from the engine, allowing it to be cleaned up. The request must have been previously added - * to the engine using OgaEngineAddRequest. After this call, the request will no longer be processed and cannot be reused. - * Removing an already closed request returns an error. + * This function logically closes a request, after which it will no longer be processed and cannot be reused. + * A nonterminal request must belong to this engine. Removing an already closed request is an engine-agnostic successful + * no-op because the request no longer has an owner. Dynamic removal releases cache ownership immediately; a resident + * static-batch row may remain physically retained until its shared batch is recycled. + * The caller remains responsible for releasing every request handle with OgaDestroyRequest. * * \param[in] engine The engine instance from which the request is being removed. - * \param[in] request The request to remove from the engine. The request must have been previously added to the engine. + * \param[in] request The request to remove. A nonterminal request must have been previously added to this engine. * \return OgaResult containing the error message if the operation failed, or nullptr on success. */ OGA_EXPORT OgaResult* OGA_API_CALL OgaEngineRemoveRequest(OgaEngine* engine, OgaRequest* request); @@ -1227,6 +1239,8 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaEngineRemoveRequest(OgaEngine* engine, Oga * * This function initializes a new request object that can be used to submit input sequences for model inference. * Once added to the engine, the request will be processed by the engine in subsequent calls to OgaEngineStep. + * The returned handle is owned by the caller. Explicit removal is recommended for deterministic resource release; + * otherwise releasing the final external handle marks a submitted request for deferred Engine reclamation. * * \param[in] params The parameters for the generator, such as temperature, top-k, etc. * \param[out] out Pointer to the created request instance. On success, *out will be set to the new request object. @@ -1262,10 +1276,12 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestContinue(OgaRequest* request, const /** * \brief Destroys the given request. * - * This function cleans up the resources associated with the request, including any input sequences and parameters. - * It should be called when the request is no longer needed, either after it has been processed. + * This function releases one external request handle. Releasing the final handle marks a submitted request abandoned; + * its Engine reclaims it before the next OgaEngineAddRequest or OgaEngineStep call. Call OgaEngineRemoveRequest first + * when resources must be released immediately. Every handle returned by OgaCreateRequest or OgaEngineStep must be + * released once. * - * \param[in] request The request to be destroyed. The request must have been created using OgaCreateRequest. + * \param[in] request A request handle returned by OgaCreateRequest or OgaEngineStep. */ OGA_EXPORT void OGA_API_CALL OgaDestroyRequest(OgaRequest* request); @@ -1327,8 +1343,21 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestGetUnseenToken(OgaRequest* request, * This function returns true at OgaRequestStatus_turn_complete. It does not mean that the request is permanently * closed; OgaRequestContinue may queue another turn while state remains resident. * - * \param[in] request The request to check if it is done. - * \param[out] out Boolean flag that will be set to true if the request is done, or false otherwise. + * \param[in] request The request whose current turn should be checked. + * \param[out] out Boolean flag that will be set to true if the current turn is complete, or false otherwise. + * \return OgaResult containing the error message if the checking of the request status failed, or nullptr on success. + */ +OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestIsTurnComplete(const OgaRequest* request, bool* out); + +/** + * \brief Deprecated compatibility alias for OgaRequestIsTurnComplete. + * + * This function reports completion of the current generation turn; it does not report permanent request closure. + * + * \deprecated Use OgaRequestIsTurnComplete instead. + * + * \param[in] request The request whose current turn should be checked. + * \param[out] out Boolean flag that will be set to true if the current turn is complete, or false otherwise. * \return OgaResult containing the error message if the checking of the request status failed, or nullptr on success. */ OGA_EXPORT OgaResult* OGA_API_CALL OgaRequestIsDone(const OgaRequest* request, bool* out); diff --git a/src/python/python.cpp b/src/python/python.cpp index c7cfe2f5d3..0758e2ec8b 100644 --- a/src/python/python.cpp +++ b/src/python/python.cpp @@ -403,6 +403,15 @@ void SetLogCallback(std::optional callback) { } } +bool IsRequestDoneDeprecated(const OgaRequest& request) { + if (PyErr_WarnEx(PyExc_DeprecationWarning, + "Request.is_done() is deprecated; use Request.is_turn_complete() instead.", + 1) < 0) { + throw pybind11::error_already_set(); + } + return request.IsTurnComplete(); +} + PYBIND11_MODULE(onnxruntime_genai, m) { m.doc() = R"pbdoc( Ort Generators library @@ -713,7 +722,7 @@ PYBIND11_MODULE(onnxruntime_genai, m) { pybind11::enum_(m, "RequestStatus") .value("CREATED", OgaRequestStatus_created) .value("QUEUED", OgaRequestStatus_queued) - .value("IN_PROGRESS", OgaRequestStatus_in_progress) + .value("ACTIVE", OgaRequestStatus_active) .value("TURN_COMPLETE", OgaRequestStatus_turn_complete) .value("CLOSED", OgaRequestStatus_closed); @@ -735,7 +744,8 @@ PYBIND11_MODULE(onnxruntime_genai, m) { request.Continue(*sequences); }) .def("has_unseen_tokens", &OgaRequest::HasUnseenTokens) - .def("is_done", &OgaRequest::IsDone) + .def("is_turn_complete", &OgaRequest::IsTurnComplete, "Return whether the current generation turn is complete.") + .def("is_done", &IsRequestDoneDeprecated, "Deprecated compatibility alias for is_turn_complete().") .def_property_readonly("status", &OgaRequest::GetStatus) .def("get_unseen_token", &OgaRequest::GetUnseenToken) .def("set_opaque_data", [](OgaRequest& request, pybind11::object opaque_data) { @@ -750,9 +760,11 @@ PYBIND11_MODULE(onnxruntime_genai, m) { pybind11::class_(m, "Engine") .def(pybind11::init([](OgaModel& model) { return OgaEngine::Create(model); })) - .def("add_request", &OgaEngine::Add) + .def("add_request", &OgaEngine::Add, + "Submit a request. The engine owns it until remove_request() is called, including after turn completion.") .def("step", &OgaEngine::Step) - .def("remove_request", &OgaEngine::Remove) + .def("remove_request", &OgaEngine::Remove, + "Remove a request. Repeated calls after it is closed are successful no-ops.") .def("has_pending_requests", &OgaEngine::HasPendingRequests); pybind11::class_(m, "StreamingProcessor") diff --git a/src/smartptrs.h b/src/smartptrs.h index 1eb308c844..a78ea947f8 100644 --- a/src/smartptrs.h +++ b/src/smartptrs.h @@ -257,15 +257,31 @@ struct DeviceInterface { // A shared_ptr based type that we expose through our C API should inherit from this type. // ExternalAddRef must be called when returning an object through the C API // ExternalRelease must be called on the C API destroy method +template +struct ExternalRefCountedTraits { + static constexpr bool notify_external_reference_changes = false; +}; + template struct ExternalRefCounted { void ExternalAddRef() { - if (++ref_count_ == 1) // First reference? + if (++ref_count_ == 1) { // First reference? external_owner_ = static_cast(this)->shared_from_this(); + if constexpr (ExternalRefCountedTraits::notify_external_reference_changes) { + static_cast(this)->OnFirstExternalReference(); + } + } } - void ExternalRelease() { - if (--ref_count_ == 0) + + void ExternalRelease() noexcept(ExternalRefCountedTraits::notify_external_reference_changes) { + if (--ref_count_ == 0) { + if constexpr (ExternalRefCountedTraits::notify_external_reference_changes) { + // Notify before releasing the self-owner so a type-specific last-release hook can only mark + // deferred work while the object is guaranteed to still be alive. + static_cast(this)->OnLastExternalReference(); + } external_owner_ = nullptr; + } } private: diff --git a/test/c_api_tests.cpp b/test/c_api_tests.cpp index 9b0111d999..5ec344bbe7 100644 --- a/test/c_api_tests.cpp +++ b/test/c_api_tests.cpp @@ -855,6 +855,26 @@ TEST(CAPITests, SetTerminate) { #endif } +TEST(CAPITests, RequestIsDoneCppCompatibilityAlias) { + auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32"); + auto params = OgaGeneratorParams::Create(*model); + auto request = OgaRequest::Create(*params); + + const bool is_turn_complete = request->IsTurnComplete(); + EXPECT_EQ(request->IsDone(), is_turn_complete); +} + +TEST(CAPITests, RequestIsDoneCCompatibilityAlias) { + auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32"); + auto params = OgaGeneratorParams::Create(*model); + auto request = OgaRequest::Create(*params); + + const bool is_turn_complete = request->IsTurnComplete(); + bool c_is_done{}; + OgaCheckResult(OgaRequestIsDone(request.get(), &c_is_done)); + EXPECT_EQ(c_is_done, is_turn_complete); +} + // DML doesn't support batch_size > 1 #if TEST_PHI2 && !USE_DML @@ -931,6 +951,13 @@ struct Phi2Test { } for (size_t i = 0; i < batch_size_; i++) { + EXPECT_EQ(requests_[i]->GetStatus(), OgaRequestStatus_turn_complete); + EXPECT_TRUE(requests_[i]->IsTurnComplete()); + EXPECT_NO_THROW(engine->Remove(*requests_[i])); + EXPECT_EQ(requests_[i]->GetStatus(), OgaRequestStatus_closed); + EXPECT_FALSE(requests_[i]->IsTurnComplete()); + EXPECT_NO_THROW(engine->Remove(*requests_[i])); + auto out_string = tokenizer_->Decode(generated_tokens[i].data(), generated_tokens[i].size()); std::cout << "Decoded string:" << out_string << std::endl; } diff --git a/test/engine/engine_invariants_tests.cpp b/test/engine/engine_invariants_tests.cpp index 36f771d75d..5e9ff01650 100644 --- a/test/engine/engine_invariants_tests.cpp +++ b/test/engine/engine_invariants_tests.cpp @@ -295,17 +295,17 @@ TEST(InvariantValidatorTest, ZeroBlockTableColumnsIsAllowed) { TEST(InvariantValidatorTest, ValidRequestHasNoViolations) { EXPECT_TRUE(ValidateRequestInvariants( - MakeValidRequest(kRequestA, RequestStatus::InProgress, 10, 4, 6)) + MakeValidRequest(kRequestA, RequestStatus::Active, 10, 4, 6)) .empty()); } TEST(InvariantValidatorTest, ProcessedBeyondCurrentReported) { - auto request = MakeValidRequest(kRequestA, RequestStatus::InProgress, 10, 12, 6); + auto request = MakeValidRequest(kRequestA, RequestStatus::Active, 10, 12, 6); EXPECT_FALSE(ValidateRequestInvariants(request).empty()); } TEST(InvariantValidatorTest, SeenBeyondCurrentReported) { - auto request = MakeValidRequest(kRequestA, RequestStatus::InProgress, 10, 4, 11); + auto request = MakeValidRequest(kRequestA, RequestStatus::Active, 10, 4, 11); EXPECT_FALSE(ValidateRequestInvariants(request).empty()); } @@ -329,8 +329,8 @@ TEST(InvariantValidatorTest, TurnCompleteRequestFullyProcessedIsValid) { TEST(InvariantValidatorTest, ConsistentSnapshotsValidateClean) { const auto cache = MakeValidCache(); const std::vector requests{ - MakeValidRequest(kRequestA, RequestStatus::InProgress, 9, 9, 9), - MakeValidRequest(kRequestB, RequestStatus::InProgress, 4, 4, 4), + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + MakeValidRequest(kRequestB, RequestStatus::Active, 4, 4, 4), }; EXPECT_TRUE(ValidateInvariants(cache, requests).empty()); EXPECT_NO_THROW(ThrowIfInvariantsViolated(cache, requests)); @@ -339,7 +339,7 @@ TEST(InvariantValidatorTest, ConsistentSnapshotsValidateClean) { TEST(InvariantValidatorTest, BlockTableForUnknownRequestReported) { const auto cache = MakeValidCache(); // owns tables for A and B const std::vector requests{ - MakeValidRequest(kRequestA, RequestStatus::InProgress, 9, 9, 9), + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), // B is missing from the request set, yet the cache holds a block table for it. }; EXPECT_FALSE(ValidateInvariants(cache, requests).empty()); @@ -349,8 +349,8 @@ TEST(InvariantValidatorTest, ThrowWrapperListsViolations) { auto cache = MakeValidCache(); cache.free_blocks = 0; // break block accounting const std::vector requests{ - MakeValidRequest(kRequestA, RequestStatus::InProgress, 9, 9, 9), - MakeValidRequest(kRequestB, RequestStatus::InProgress, 4, 4, 4), + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + MakeValidRequest(kRequestB, RequestStatus::Active, 4, 4, 4), }; EXPECT_THROW(ThrowIfInvariantsViolated(cache, requests), std::runtime_error); } diff --git a/test/engine/engine_step_tests.cpp b/test/engine/engine_step_tests.cpp index 8d277897c5..9ed65a6049 100644 --- a/test/engine/engine_step_tests.cpp +++ b/test/engine/engine_step_tests.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -32,6 +33,32 @@ int IndexOf(const CallTrace& trace, const std::string& entry) { return it == trace.entries.end() ? -1 : static_cast(it - trace.entries.begin()); } +class ExternalRequestReference { + public: + explicit ExternalRequestReference(Request& request) : request_{&request} { + request_->ExternalAddRef(); + } + + ExternalRequestReference(const ExternalRequestReference&) = delete; + ExternalRequestReference& operator=(const ExternalRequestReference&) = delete; + + ~ExternalRequestReference() { + Release(); + } + + void Release() noexcept { + if (request_) { + request_->ExternalRelease(); + request_ = nullptr; + } + } + + private: + Request* request_; +}; + +static_assert(noexcept(std::declval().ExternalRelease())); + class EngineStepTest : public ::testing::Test { protected: void SetUp() override { @@ -57,7 +84,7 @@ TEST_F(EngineStepTest, SingleRequestSchedulesThenDecodesThenReturns) { ASSERT_NE(ready, nullptr); EXPECT_EQ(ready, request); - EXPECT_TRUE(request->IsDone()); + EXPECT_TRUE(request->IsTurnComplete()); EXPECT_EQ(engine.executor->decode_calls, 1); ASSERT_EQ(engine.executor->decoded_batch_sizes.size(), 1u); EXPECT_EQ(engine.executor->decoded_batch_sizes[0], 1u); @@ -74,6 +101,8 @@ TEST_F(EngineStepTest, SingleRequestSchedulesThenDecodesThenReturns) { TEST_F(EngineStepTest, FittingRequestsShareOneDecodeAndDrainWithoutReexecuting) { auto engine = MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); + // These test-only shared_ptr requests never acquire an external reference. Repeated Step + // boundaries must not mistake them for abandoned public handles. std::vector> requests; for (int32_t seed : {10, 20, 30}) { auto prompt = Prompt(seed); @@ -93,7 +122,7 @@ TEST_F(EngineStepTest, FittingRequestsShareOneDecodeAndDrainWithoutReexecuting) std::sort(returned.begin(), returned.end()); EXPECT_EQ(returned, sorted_requests); for (const auto& request : requests) { - EXPECT_TRUE(request->IsDone()); + EXPECT_TRUE(request->IsTurnComplete()); } EXPECT_EQ(engine.executor->decode_calls, 1); ASSERT_EQ(engine.executor->decoded_batch_sizes.size(), 1u); @@ -101,15 +130,173 @@ TEST_F(EngineStepTest, FittingRequestsShareOneDecodeAndDrainWithoutReexecuting) EXPECT_FALSE(engine.engine->HasPendingRequests()); } +TEST_F(EngineStepTest, AddRequestReclaimsAbandonedTurnCompleteAtCapacity) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/1, EosToken(*model_)); + auto first_prompt = Prompt(10); + auto first = MintRequest(*model_, first_prompt); + ExternalRequestReference first_external{*first}; + engine.engine->AddRequest(first); + + ASSERT_EQ(engine.engine->Step(), first); + ASSERT_EQ(first->status_, RequestStatus::TurnComplete); + ASSERT_EQ(engine.cache->AllocatedCount(), 1u); + + first_external.Release(); + + auto second_prompt = Prompt(20); + auto second = MintRequest(*model_, second_prompt); + ExternalRequestReference second_external{*second}; + engine.engine->AddRequest(second); + + EXPECT_EQ(first->status_, RequestStatus::Closed); + EXPECT_EQ(engine.cache->AllocatedCount(), 0u); + EXPECT_EQ(engine.cache->deallocate_calls, 1); + + EXPECT_EQ(engine.engine->Step(), second); + EXPECT_EQ(second->status_, RequestStatus::TurnComplete); + EXPECT_EQ(engine.cache->AllocatedCount(), 1u); + + engine.engine->RemoveRequest(second); + second_external.Release(); +} + +TEST_F(EngineStepTest, StepReclaimsAbandonedTurnCompleteBeforePlanningAtCapacity) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/1, EosToken(*model_)); + auto first_prompt = Prompt(10); + auto second_prompt = Prompt(20); + auto first = MintRequest(*model_, first_prompt); + auto second = MintRequest(*model_, second_prompt); + ExternalRequestReference first_external{*first}; + ExternalRequestReference second_external{*second}; + engine.engine->AddRequest(first); + engine.engine->AddRequest(second); + + ASSERT_EQ(engine.engine->Step(), first); + ASSERT_EQ(first->status_, RequestStatus::TurnComplete); + ASSERT_EQ(second->status_, RequestStatus::Assigned); + ASSERT_EQ(engine.cache->AllocatedCount(), 1u); + first_external.Release(); + + EXPECT_EQ(engine.engine->Step(), second); + EXPECT_EQ(first->status_, RequestStatus::Closed); + EXPECT_EQ(second->status_, RequestStatus::TurnComplete); + EXPECT_EQ(engine.cache->AllocatedCount(), 1u); + EXPECT_EQ(engine.cache->deallocate_calls, 1); + EXPECT_EQ(engine.executor->decode_calls, 2); + + engine.engine->RemoveRequest(second); + second_external.Release(); +} + +TEST_F(EngineStepTest, StepPurgesAbandonedReadyAndQueuedRequestsExactlyOnce) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/2, EosToken(*model_)); + auto survivor_prompt = Prompt(10); + auto ready_orphan_prompt = Prompt(20); + auto queued_orphan_prompt = Prompt(30); + auto survivor = MintRequest(*model_, survivor_prompt); + auto ready_orphan = MintRequest(*model_, ready_orphan_prompt); + auto queued_orphan = MintRequest(*model_, queued_orphan_prompt); + ExternalRequestReference survivor_external{*survivor}; + ExternalRequestReference ready_orphan_external{*ready_orphan}; + ExternalRequestReference queued_orphan_external{*queued_orphan}; + engine.engine->AddRequest(survivor); + engine.engine->AddRequest(ready_orphan); + engine.engine->AddRequest(queued_orphan); + + ASSERT_EQ(engine.engine->Step(), survivor); + ASSERT_EQ(survivor->status_, RequestStatus::TurnComplete); + ASSERT_EQ(ready_orphan->status_, RequestStatus::TurnComplete); + ASSERT_EQ(queued_orphan->status_, RequestStatus::Assigned); + ASSERT_EQ(engine.cache->AllocatedCount(), 2u); + + ready_orphan_external.Release(); + queued_orphan_external.Release(); + ASSERT_EQ(engine.cache->deallocate_calls, 0); + + // Cleanup runs before Step can drain the orphan's ready notification or plan the queued orphan. + EXPECT_EQ(engine.engine->Step(), nullptr); + EXPECT_EQ(ready_orphan->status_, RequestStatus::Closed); + EXPECT_EQ(queued_orphan->status_, RequestStatus::Closed); + EXPECT_EQ(survivor->status_, RequestStatus::TurnComplete); + EXPECT_EQ(engine.cache->AllocatedCount(), 1u); + EXPECT_EQ(engine.cache->deallocate_calls, 2); + EXPECT_EQ(engine.executor->decode_calls, 1); + + // Closed requests were removed from Engine ownership, so another boundary cannot clean them twice. + EXPECT_EQ(engine.engine->Step(), nullptr); + EXPECT_EQ(engine.cache->deallocate_calls, 2); + + engine.engine->RemoveRequest(survivor); + survivor_external.Release(); +} + +TEST_F(EngineStepTest, ReacquiringExternalReferenceCancelsDeferredAbandonment) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/1, EosToken(*model_)); + auto prompt = Prompt(10); + auto request = MintRequest(*model_, prompt); + ExternalRequestReference initial_external{*request}; + engine.engine->AddRequest(request); + ASSERT_EQ(engine.engine->Step(), request); + ASSERT_EQ(request->status_, RequestStatus::TurnComplete); + + initial_external.Release(); + ExternalRequestReference reacquired_external{*request}; + + EXPECT_EQ(engine.engine->Step(), nullptr); + EXPECT_EQ(request->status_, RequestStatus::TurnComplete); + EXPECT_EQ(engine.cache->AllocatedCount(), 1u); + EXPECT_EQ(engine.cache->deallocate_calls, 0); + + engine.engine->RemoveRequest(request); + reacquired_external.Release(); +} + +TEST_F(EngineStepTest, ContinueRejectsUndrainedReadyNotificationWithoutMutation) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); + auto first_prompt = Prompt(10); + auto second_prompt = Prompt(20); + auto first = MintRequest(*model_, first_prompt); + auto second = MintRequest(*model_, second_prompt); + engine.engine->AddRequest(first); + engine.engine->AddRequest(second); + + ASSERT_EQ(engine.engine->Step(), first); + ASSERT_EQ(engine.executor->decode_calls, 1); + ASSERT_EQ(first->status_, RequestStatus::TurnComplete); + ASSERT_EQ(second->status_, RequestStatus::TurnComplete); + + const auto before = second->Snapshot(); + const std::vector continuation{5, 6}; + try { + second->Continue(continuation); + FAIL() << "Expected Continue to reject the undrained ready notification."; + } catch (const std::runtime_error& error) { + EXPECT_NE(std::string{error.what()}.find("Engine::Step()"), std::string::npos); + } + + const auto rejected = second->Snapshot(); + EXPECT_EQ(rejected.status, before.status); + EXPECT_EQ(rejected.current_sequence_length, before.current_sequence_length); + EXPECT_EQ(rejected.processed_sequence_length, before.processed_sequence_length); + + EXPECT_EQ(engine.engine->Step(), second); + EXPECT_EQ(engine.executor->decode_calls, 1); + + EXPECT_NO_THROW(second->Continue(continuation)); + const auto continued = second->Snapshot(); + EXPECT_EQ(continued.status, RequestStatus::Assigned); + EXPECT_EQ(continued.current_sequence_length, + before.current_sequence_length + static_cast(continuation.size())); +} + // Under capacity backpressure Step decodes only the requests that fit, then forms a fresh batch for // the deferred request on a later step -- one decode per internal cycle, never an over-capacity run. // -// Cache residency is released only by an explicit RemoveRequest, so a finished request keeps its -// slot until the caller gives it back. With capacity for two requests, the third is admitted only -// after a completed request has been removed. The loop therefore removes each ready request as the -// caller is expected to, which is what frees the slot the deferred request needs; the assertions -// pin both halves of that contract -- the slot is still held when the request is handed back, and -// it is released exactly at removal. +// These test requests use only internal shared_ptrs and have never had public handles, so they are +// not abandoned automatically. A finished request keeps its slot until explicitly removed. With +// capacity for two requests, the third is admitted only after a completed request has been removed. +// The assertions pin both halves of that contract -- the slot is still held when the request is +// handed back, and it is released exactly at removal. TEST_F(EngineStepTest, BackpressureFormsAFreshBatchAcrossSteps) { auto engine = MakeDoublesEngine(model_, /*capacity=*/2, EosToken(*model_)); @@ -125,7 +312,7 @@ TEST_F(EngineStepTest, BackpressureFormsAFreshBatchAcrossSteps) { int removals = 0; while (auto ready = engine.engine->Step()) { // The fixture's executor forces end-of-stream, so every request Step hands back is finished. - EXPECT_TRUE(ready->IsDone()); + EXPECT_TRUE(ready->IsTurnComplete()); // The finished request still owns its cache slot: nothing is reclaimed implicitly. EXPECT_EQ(engine.cache->deallocate_calls, removals); returned.push_back(ready); @@ -190,6 +377,38 @@ TEST_F(EngineStepTest, StaticBatchingPreservesOrderingAndReusesResidentContinuat EXPECT_EQ(cache_observer->allocate_calls, allocations_before); } +TEST_F(EngineStepTest, StaticStepClosesAbandonedResidentWithoutIndividualDeallocation) { + model_->config_->engine.dynamic_batching.reset(); + auto cache = std::make_shared( + model_, /*capacity=*/1, nullptr, /*supports_dynamic_batching=*/false); + auto scheduler = Scheduler::Create(model_, cache); + auto executor = std::make_unique( + model_, cache, EosToken(*model_)); + auto* executor_observer = executor.get(); + EngineDependencies dependencies{cache, std::move(scheduler), + std::move(executor)}; + auto engine = std::make_shared(model_, std::move(dependencies)); + auto prompt = Prompt(10); + auto request = MintRequest(*model_, prompt); + ExternalRequestReference external{*request}; + engine->AddRequest(request); + + ASSERT_EQ(engine->Step(), request); + ASSERT_EQ(request->status_, RequestStatus::TurnComplete); + ASSERT_EQ(cache->AllocatedCount(), 1u); + external.Release(); + + EXPECT_EQ(engine->Step(), nullptr); + EXPECT_EQ(request->status_, RequestStatus::Closed); + EXPECT_EQ(executor_observer->decode_calls, 1); + EXPECT_EQ(cache->AllocatedCount(), 1u); + EXPECT_EQ(cache->deallocate_calls, 0); + + // A static row is logically closed once but remains physically resident until batch recycling. + EXPECT_EQ(engine->Step(), nullptr); + EXPECT_EQ(cache->deallocate_calls, 0); +} + TEST_F(EngineStepTest, StaticContinueFailsAfterBatchRecycling) { model_->config_->engine.dynamic_batching.reset(); auto cache = std::make_shared(model_); @@ -289,7 +508,7 @@ TEST_F(EngineStepTest, RetryableExecutionFailureRollsBackAndCanRetry) { auto ready = engine.engine->Step(); EXPECT_EQ(ready, request); - EXPECT_TRUE(request->IsDone()); + EXPECT_TRUE(request->IsTurnComplete()); } TEST_F(EngineStepTest, ContinuedResidentRollsBackToQueuedAndCanRetry) { @@ -402,7 +621,7 @@ TEST_F(EngineStepTest, ExecutionCapacityFailureRollsBackWithoutPoisoningEngine) auto ready = engine.engine->Step(); EXPECT_EQ(ready, request); - EXPECT_TRUE(request->IsDone()); + EXPECT_TRUE(request->IsTurnComplete()); } TEST_F(EngineStepTest, PostProcessingFailureRestoresSearchAndCanRetry) { @@ -432,7 +651,7 @@ TEST_F(EngineStepTest, PostProcessingFailureRestoresSearchAndCanRetry) { auto ready = engine.engine->Step(); EXPECT_EQ(ready, request); - EXPECT_TRUE(request->IsDone()); + EXPECT_TRUE(request->IsTurnComplete()); } TEST_F(EngineStepTest, LaterRequestFailureRestoresEarlierSample) { diff --git a/test/engine/request_lifecycle_tests.cpp b/test/engine/request_lifecycle_tests.cpp index 6680aaabb9..e0d8dd33e0 100644 --- a/test/engine/request_lifecycle_tests.cpp +++ b/test/engine/request_lifecycle_tests.cpp @@ -4,7 +4,7 @@ // Lifecycle tests for the engine Request state machine. Because a tiny real // CPU fixture model is available, these tests drive genuine Request objects (rather than a mock // Search) and pin the transition policy: which mutations each status permits, and how -// create/assign/schedule/continue/remove move a request between Unassigned, Assigned, InProgress, +// create/assign/schedule/continue/remove move a request between Unassigned, Assigned, Active, // TurnComplete, and Closed. #include @@ -107,26 +107,26 @@ TEST_F(RequestLifecycleTest, ScheduleIsRejectedBeforeAssign) { EXPECT_EQ(request->status_, RequestStatus::Unassigned); } -// An assigned, non-empty request schedules cleanly and moves to InProgress. -TEST_F(RequestLifecycleTest, ScheduleFromAssignedMovesToInProgress) { +// An assigned, non-empty request schedules cleanly and moves to Active. +TEST_F(RequestLifecycleTest, ScheduleFromAssignedMovesToActive) { auto prompt = Prompt(); auto request = MintAssignedRequest(engine_.engine, *model_, prompt); request->Schedule(); - EXPECT_EQ(request->status_, RequestStatus::InProgress); + EXPECT_EQ(request->status_, RequestStatus::Active); } -// While a request is in progress its token stream is owned by the engine, so external appends are +// While a request is active its token stream is owned by the engine, so external appends are // rejected without mutating the request. -TEST_F(RequestLifecycleTest, AppendIsRejectedWhileInProgress) { +TEST_F(RequestLifecycleTest, AppendIsRejectedWhileActive) { auto prompt = Prompt(); auto request = MintAssignedRequest(engine_.engine, *model_, prompt); request->Schedule(); - ASSERT_EQ(request->status_, RequestStatus::InProgress); + ASSERT_EQ(request->status_, RequestStatus::Active); const int64_t length_before = request->CurrentSequenceLength(); std::vector more{5, 6}; EXPECT_THROW(request->AddTokens(more), std::runtime_error); - EXPECT_EQ(request->status_, RequestStatus::InProgress); + EXPECT_EQ(request->status_, RequestStatus::Active); EXPECT_EQ(request->CurrentSequenceLength(), length_before); } @@ -153,9 +153,11 @@ TEST_F(RequestLifecycleTest, ContinueAfterTurnCompleteQueuesNextTurn) { EXPECT_EQ(request->CurrentSequenceLength(), assigned_length + static_cast(more.size())); EXPECT_EQ(request->status_, RequestStatus::Assigned); - EXPECT_FALSE(request->IsDone()); + EXPECT_FALSE(request->IsTurnComplete()); EXPECT_EQ(engine_.engine->Step(), request); EXPECT_EQ(request->status_, RequestStatus::TurnComplete); + EXPECT_TRUE(request->IsTurnComplete()); + EXPECT_TRUE(request->IsDone()); // Compatibility alias. } TEST_F(RequestLifecycleTest, ContinueBeyondContextIsRejectedBeforeMutation) { @@ -236,7 +238,7 @@ TEST_F(RequestLifecycleTest, ContinueIsRejectedOutsideTurnComplete) { } // Removing a request releases it from the engine and makes it terminal. -TEST_F(RequestLifecycleTest, RemoveMakesRequestTerminal) { +TEST_F(RequestLifecycleTest, RequestRemoveIsIdempotentAfterClose) { auto prompt = Prompt(); const std::vector more{5}; auto request = MintAssignedRequest(engine_.engine, *model_, prompt); @@ -248,7 +250,36 @@ TEST_F(RequestLifecycleTest, RemoveMakesRequestTerminal) { EXPECT_EQ(engine_.cache->AllocatedCount(), 0u); EXPECT_THROW(request->AddTokens(more), std::runtime_error); EXPECT_THROW(request->Continue(more), std::runtime_error); - EXPECT_THROW(request->Remove(), std::runtime_error); + EXPECT_NO_THROW(request->Remove()); + EXPECT_EQ(engine_.cache->deallocate_calls, 1); +} + +TEST_F(RequestLifecycleTest, EngineRemoveRequestIsIdempotentAfterClose) { + auto other_engine = + MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); + auto prompt = Prompt(); + auto request = MintAssignedRequest(engine_.engine, *model_, prompt); + + engine_.engine->RemoveRequest(request); + ASSERT_EQ(request->status_, RequestStatus::Closed); + ASSERT_EQ(engine_.cache->deallocate_calls, 1); + + EXPECT_NO_THROW(engine_.engine->RemoveRequest(request)); + EXPECT_NO_THROW(other_engine.engine->RemoveRequest(request)); + EXPECT_EQ(engine_.cache->deallocate_calls, 1); + EXPECT_EQ(other_engine.cache->deallocate_calls, 0); +} + +TEST_F(RequestLifecycleTest, EngineRemoveRequestRejectsNonterminalRequestFromAnotherEngine) { + auto other_engine = + MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); + auto prompt = Prompt(); + auto request = MintAssignedRequest(engine_.engine, *model_, prompt); + + EXPECT_THROW(other_engine.engine->RemoveRequest(request), std::runtime_error); + EXPECT_EQ(request->status_, RequestStatus::Assigned); + EXPECT_EQ(engine_.cache->deallocate_calls, 0); + EXPECT_EQ(other_engine.cache->deallocate_calls, 0); } TEST_F(RequestLifecycleTest, RemoveIsRejectedBeforeSubmission) { @@ -296,7 +327,7 @@ TEST_F(RequestLifecycleTest, TransactionalLogitsStageUntilCommit) { request->CommitStep(plan, result); const auto committed = request->Snapshot(); - EXPECT_EQ(committed.status, RequestStatus::InProgress); + EXPECT_EQ(committed.status, RequestStatus::Active); EXPECT_EQ(committed.processed_sequence_length, before.current_sequence_length); EXPECT_FALSE(committed.is_prefill); ASSERT_TRUE(request->HasUnseenTokens()); @@ -337,7 +368,7 @@ TEST_F(RequestLifecycleTest, PartialPrefillAdvancesOnlyAtCommit) { request->CommitStep(plan, RequestStepResult{}); const auto committed = request->Snapshot(); - EXPECT_EQ(committed.status, RequestStatus::InProgress); + EXPECT_EQ(committed.status, RequestStatus::Active); EXPECT_EQ(committed.current_sequence_length, before.current_sequence_length); EXPECT_EQ(committed.processed_sequence_length, 2); // Two of the three prompt tokens are in the cache, so the request is still prefilling. diff --git a/test/engine/scheduler_contract_tests.cpp b/test/engine/scheduler_contract_tests.cpp index 3f522e33d9..52a5523b58 100644 --- a/test/engine/scheduler_contract_tests.cpp +++ b/test/engine/scheduler_contract_tests.cpp @@ -277,7 +277,7 @@ TEST_F(SchedulerContractTest, DynamicPlanningKeepsActiveWorkWhenNewAdmissionIsDe DynamicBatchScheduler scheduler(model_, cache); auto active = Assigned(10); MakePrefillResident(scheduler, *cache, active); - ASSERT_EQ(active->status_, RequestStatus::InProgress); + ASSERT_EQ(active->status_, RequestStatus::Active); auto deferred = Assigned(20); scheduler.AddRequest(deferred); diff --git a/test/python/test_onnxruntime_genai_engine.py b/test/python/test_onnxruntime_genai_engine.py index 2114e41437..a8b192bc14 100644 --- a/test/python/test_onnxruntime_genai_engine.py +++ b/test/python/test_onnxruntime_genai_engine.py @@ -63,6 +63,22 @@ def model(device): return og.Model(config) +def test_request_status_names(): + assert hasattr(og.RequestStatus, "ACTIVE") + assert not hasattr(og.RequestStatus, "IN_PROGRESS") + + +def test_request_is_done_compatibility_alias_warns(model): + params = og.GeneratorParams(model) + request = og.Request(params) + expected = request.is_turn_complete() + + with pytest.warns(DeprecationWarning, match=r"is_done.*is_turn_complete"): + actual = request.is_done() + + assert actual == expected + + def _add_request(engine, model, prompt, max_new_tokens, sink): params = og.GeneratorParams(model) params.set_search_options(do_sample=False, max_length=len(prompt) + max_new_tokens) @@ -77,7 +93,7 @@ def _drain(ready): sink = ready.get_opaque_data() while ready.has_unseen_tokens(): sink.tokens.append(ready.get_unseen_token()) - return ready.is_done() + return ready.is_turn_complete() def _step_once(engine): @@ -231,7 +247,7 @@ def test_continuation_while_peer_remains_active(model): reference = _add_request( reference_engine, model, _PROMPT_A, short_max_new, reference_sink ) - while not reference.is_done(): + while not reference.is_turn_complete(): ready = reference_engine.step() assert ready is not None _drain(ready) @@ -243,18 +259,18 @@ def test_continuation_while_peer_remains_active(model): short = _add_request(engine, model, _PROMPT_A, short_max_new, short_sink) long = _add_request(engine, model, _PROMPT_LONG, long_max_new, long_sink) - while not short.is_done(): + while not short.is_turn_complete(): ready = engine.step() assert ready is not None if _drain(ready) and ready is not short: engine.remove_request(ready) - assert not long.is_done(), "peer must remain active when continuation is appended" + assert not long.is_turn_complete(), "peer must remain active when continuation is appended" for _ in range(3): ready = engine.step() assert ready is not None _drain(ready) - assert not long.is_done(), "peer must remain active during the continuation delay" + assert not long.is_turn_complete(), "peer must remain active during the continuation delay" short.continue_with(np.asarray(follow_up, dtype=np.int32)) _run(engine) @@ -283,6 +299,8 @@ def test_request_cannot_be_removed_from_another_engine(model): other.remove_request(request) owner.remove_request(request) + other.remove_request(request) + assert request.status == og.RequestStatus.CLOSED def test_request_lifecycle_status(model): @@ -298,11 +316,18 @@ def test_request_lifecycle_status(model): engine.add_request(request) assert request.status == og.RequestStatus.QUEUED - while not request.is_done(): + ready = engine.step() + assert ready is not None + _drain(ready) + assert request.status == og.RequestStatus.ACTIVE + assert not request.is_turn_complete() + + while not request.is_turn_complete(): ready = engine.step() assert ready is not None _drain(ready) assert request.status == og.RequestStatus.TURN_COMPLETE + assert request.is_turn_complete() with pytest.raises(RuntimeError, match="use Continue"): request.add_tokens(np.asarray([12], dtype=np.int32)) @@ -311,12 +336,13 @@ def test_request_lifecycle_status(model): engine.remove_request(request) assert request.status == og.RequestStatus.CLOSED + assert not request.is_turn_complete() with pytest.raises(RuntimeError, match="closed request"): request.add_tokens(np.asarray([12], dtype=np.int32)) with pytest.raises(RuntimeError, match="closed request"): request.continue_with(np.asarray([12], dtype=np.int32)) - with pytest.raises(RuntimeError, match="already closed"): - engine.remove_request(request) + engine.remove_request(request) + assert request.status == og.RequestStatus.CLOSED def test_remove_request_freezes_output(model): From 5f7fa8d1c6242cde7d566c0e3605b6eac895d6be Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 19:12:24 -0500 Subject: [PATCH 06/19] Add windowed continuation wrap and rollback coverage Introduce a deterministic two-layer synthetic paged model that exercises two sliding-window ring wraps and verifies exact continuation-versus-replay output. Run a real windowed decode before injecting a retryable failure, then verify cache ownership, request state, reservation cleanup, and retry output. Files: .gitignore; test/models/engine/synthetic-windowed-multiwrap/*; test/python/create/create_synthetic_windowed_multiwrap_model.py; test/python/models/test_engine_windowed_multiwrap.py; test/engine/windowed_transaction_tests.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b232f82c-f25c-429a-a855-7f8c8f40bbf3 --- .gitignore | 6 + test/engine/windowed_transaction_tests.cpp | 385 +++++++++++ .../synthetic-windowed-multiwrap/decoder.onnx | Bin 0 -> 10403 bytes .../genai_config.json | 56 ++ ...eate_synthetic_windowed_multiwrap_model.py | 631 ++++++++++++++++++ .../models/test_engine_windowed_multiwrap.py | 199 ++++++ 6 files changed, 1277 insertions(+) create mode 100644 test/engine/windowed_transaction_tests.cpp create mode 100644 test/models/engine/synthetic-windowed-multiwrap/decoder.onnx create mode 100644 test/models/engine/synthetic-windowed-multiwrap/genai_config.json create mode 100644 test/python/create/create_synthetic_windowed_multiwrap_model.py create mode 100644 test/python/models/test_engine_windowed_multiwrap.py diff --git a/.gitignore b/.gitignore index ec7a8ed488..bc70ee352d 100644 --- a/.gitignore +++ b/.gitignore @@ -44,6 +44,12 @@ examples/csharp/ModelChat/models !test/models/qwen3-5/* !test/models/qwen3-vl/* !test/models/whisper/* +!/test/models/engine/ +/test/models/engine/* +!/test/models/engine/synthetic-windowed-multiwrap/ +/test/models/engine/synthetic-windowed-multiwrap/* +!/test/models/engine/synthetic-windowed-multiwrap/decoder.onnx +!/test/models/engine/synthetic-windowed-multiwrap/genai_config.json .ipynb_checkpoints/ /src/java/.gradle diff --git a/test/engine/windowed_transaction_tests.cpp b/test/engine/windowed_transaction_tests.cpp new file mode 100644 index 0000000000..4c6805b08f --- /dev/null +++ b/test/engine/windowed_transaction_tests.cpp @@ -0,0 +1,385 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "engine/cache_manager.h" +#include "engine/engine.h" +#include "engine/engine_invariants.h" +#include "engine/model_executor.h" +#include "engine/scheduler.h" +#include "engine_test_helpers.h" + +namespace Generators { +namespace test { +namespace { + +constexpr int32_t kInvariantFailureToken = 31; +const std::vector kInitialPrompt{4, 6}; +const std::vector kFirstTurnOutput{12}; +const std::vector kContinuation{8, 3, 10, 5, 12, 7, 14, 9, 16}; +const std::vector kSecondTurnOutput{28, 17}; + +std::shared_ptr LoadWindowedMultiwrapModel() { + return CreateModel( + GetOrtEnv(), MODEL_PATH "engine/synthetic-windowed-multiwrap"); +} + +std::vector RunTurn(const std::shared_ptr& engine, + const std::shared_ptr& request) { + std::vector output; + size_t steps = 0; + while (!request->IsTurnComplete()) { + if (++steps >= 100) { + throw std::runtime_error( + "Synthetic windowed request did not complete."); + } + + auto ready = engine->Step(); + if (!ready) { + continue; + } + if (ready != request) { + throw std::runtime_error( + "Synthetic windowed Engine returned an unexpected request."); + } + while (ready->HasUnseenTokens()) { + output.push_back(ready->UnseenToken()); + } + } + return output; +} + +// Wraps the real paged cache only to make reservation lifetime observable. All +// planning, block-table construction, cache binding, and ownership changes +// still run through PagedCacheManager. +class ObservingPagedCacheManager final : public PagedCacheManager { + public: + explicit ObservingPagedCacheManager(std::shared_ptr model) + : PagedCacheManager(std::move(model)) {} + + std::unique_ptr ReserveStep( + const StepPlan& plan) override { + class ObservedReservation final : public CacheStepReservation { + public: + ObservedReservation( + ObservingPagedCacheManager& owner, + std::unique_ptr inner) + : owner_{owner}, inner_{std::move(inner)} {} + + ~ObservedReservation() override { + if (active_) { + active_ = false; + owner_.OnAbandonedReservation(); + } + } + + PagedCacheReservation* PagedReservation() override { + return inner_->PagedReservation(); + } + + void Commit() override { + inner_->Commit(); + if (active_) { + active_ = false; + owner_.OnCommittedReservation(); + } + } + + void Release() override { + inner_->Release(); + if (active_) { + active_ = false; + owner_.OnReleasedReservation(); + } + } + + private: + ObservingPagedCacheManager& owner_; + std::unique_ptr inner_; + bool active_{true}; + }; + + auto reservation = PagedCacheManager::ReserveStep(plan); + ++reservation_count_; + ++active_reservations_; + return std::make_unique( + *this, std::move(reservation)); + } + + size_t ReservationCount() const { return reservation_count_; } + size_t ActiveReservations() const { return active_reservations_; } + size_t CommitCount() const { return commit_count_; } + size_t ReleaseCount() const { return release_count_; } + size_t AbandonedReservationCount() const { + return abandoned_reservation_count_; + } + + private: + void OnCommittedReservation() { + --active_reservations_; + ++commit_count_; + } + + void OnReleasedReservation() { + --active_reservations_; + ++release_count_; + } + + void OnAbandonedReservation() { + --active_reservations_; + ++abandoned_reservation_count_; + } + + size_t reservation_count_{}; + size_t active_reservations_{}; + size_t commit_count_{}; + size_t release_count_{}; + size_t abandoned_reservation_count_{}; +}; + +// Delegates to the production DecoderModelExecutor first. Returning from that +// delegate means the synchronous ORT run completed with the paged cache bound +// as both input and output. Only then does this wrapper observe the real logits +// and raise a retryable failure. +class RetryableAfterRealDecodeExecutor final : public ModelExecutor { + public: + explicit RetryableAfterRealDecodeExecutor( + std::unique_ptr inner) + : inner_{std::move(inner)} {} + + void Decode(ScheduledRequests& scheduled_requests, + ExecutionContext& context) override { + inner_->Decode(scheduled_requests, context); + ++completed_real_decodes_; + + if (!fail_next_decode_) { + return; + } + fail_next_decode_ = false; + + const auto logits = scheduled_requests.ProcessLogits(); + observed_real_logits_ = logits.size() == scheduled_requests.size(); + if (!context.plan || context.plan->requests.size() != 1) { + throw std::logic_error( + "Windowed rollback fault expected one planned request."); + } + const auto& entry = context.plan->requests.front(); + failed_unprocessed_token_count_ = entry.unprocessed_token_count; + failed_target_cache_slots_ = entry.target_cache_slots; + injected_failure_ = true; + throw ModelExecutionError{ + ExecutionFailureKind::RetryableAbort, + "Injected retryable failure after real windowed model execution.", + }; + } + + void FailNextDecode() { fail_next_decode_ = true; } + + bool InjectedFailure() const { return injected_failure_; } + bool ObservedRealLogits() const { return observed_real_logits_; } + size_t CompletedRealDecodes() const { return completed_real_decodes_; } + size_t FailedUnprocessedTokenCount() const { + return failed_unprocessed_token_count_; + } + size_t FailedTargetCacheSlots() const { + return failed_target_cache_slots_; + } + + private: + std::unique_ptr inner_; + bool fail_next_decode_{}; + bool injected_failure_{}; + bool observed_real_logits_{}; + size_t completed_real_decodes_{}; + size_t failed_unprocessed_token_count_{}; + size_t failed_target_cache_slots_{}; +}; + +struct FaultInjectingEngine { + std::shared_ptr engine; + std::shared_ptr cache; + RetryableAfterRealDecodeExecutor* executor{}; +}; + +FaultInjectingEngine MakeFaultInjectingEngine( + const std::shared_ptr& model) { + auto cache = + std::make_shared(model); + auto scheduler = Scheduler::Create(model, cache); + auto executor = std::make_unique( + ModelExecutor::Create(model, cache)); + auto* executor_observer = executor.get(); + EngineDependencies dependencies{ + cache, std::move(scheduler), std::move(executor)}; + auto engine = + std::make_shared(model, std::move(dependencies)); + return FaultInjectingEngine{ + std::move(engine), std::move(cache), executor_observer}; +} + +void ExpectRequestBlocksEqual(const RequestBlockSnapshot& actual, + const RequestBlockSnapshot& expected) { + EXPECT_EQ(actual.request_id, expected.request_id); + EXPECT_EQ(actual.block_ids, expected.block_ids); + EXPECT_EQ(actual.used_slots, expected.used_slots); + EXPECT_EQ(actual.empty_slots, expected.empty_slots); +} + +void ExpectCacheOwnershipRestored(const PagedCacheSnapshot& actual, + const PagedCacheSnapshot& expected) { + EXPECT_EQ(actual.block_size, expected.block_size); + EXPECT_EQ(actual.total_blocks, expected.total_blocks); + EXPECT_EQ(actual.free_blocks, expected.free_blocks); + EXPECT_EQ(actual.AllocatedBlocks(), expected.AllocatedBlocks()); + EXPECT_TRUE(actual.transaction_reserved_block_ids.empty()); + EXPECT_TRUE(actual.reservations.empty()); + + ASSERT_EQ(actual.requests.size(), expected.requests.size()); + for (size_t i = 0; i < actual.requests.size(); ++i) { + ExpectRequestBlocksEqual(actual.requests[i], expected.requests[i]); + } + + EXPECT_EQ(actual.window_blocks.total_blocks, + expected.window_blocks.total_blocks); + EXPECT_EQ(actual.window_blocks.free_blocks, + expected.window_blocks.free_blocks); + EXPECT_EQ(actual.window_blocks.blocks_per_request, + expected.window_blocks.blocks_per_request); + EXPECT_TRUE( + actual.window_blocks.transaction_reserved_block_ids.empty()); + ASSERT_EQ(actual.window_blocks.requests.size(), + expected.window_blocks.requests.size()); + for (size_t i = 0; i < actual.window_blocks.requests.size(); ++i) { + ExpectRequestBlocksEqual(actual.window_blocks.requests[i], + expected.window_blocks.requests[i]); + } +} + +std::vector RunCleanReplay() { + auto model = LoadWindowedMultiwrapModel(); + auto engine = std::make_shared(model); + std::vector replay_prompt = kInitialPrompt; + replay_prompt.insert(replay_prompt.end(), kFirstTurnOutput.begin(), + kFirstTurnOutput.end()); + replay_prompt.insert(replay_prompt.end(), kContinuation.begin(), + kContinuation.end()); + auto request = MintRequest(*model, replay_prompt); + engine->AddRequest(request); + + auto output = RunTurn(engine, request); + engine->RemoveRequest(request); + return output; +} + +TEST(WindowedTransactionTest, + ContinuedRingWritesRollbackAndRetryMatchCleanReplay) { + const auto clean_replay_output = RunCleanReplay(); + ASSERT_EQ(clean_replay_output, kSecondTurnOutput); + ASSERT_EQ(std::find(clean_replay_output.begin(), + clean_replay_output.end(), + kInvariantFailureToken), + clean_replay_output.end()); + + auto model = LoadWindowedMultiwrapModel(); + auto faulting = MakeFaultInjectingEngine(model); + auto request = MintRequest(*model, kInitialPrompt); + faulting.engine->AddRequest(request); + ASSERT_EQ(RunTurn(faulting.engine, request), kFirstTurnOutput); + ASSERT_EQ(request->Status(), RequestStatus::TurnComplete); + ASSERT_FALSE(request->HasUnseenTokens()); + + request->Continue(kContinuation); + const auto request_before = request->Snapshot(); + const auto cache_before = faulting.cache->Snapshot(); + ASSERT_EQ(request_before.status, RequestStatus::Assigned); + ASSERT_EQ(request_before.current_sequence_length, 12); + ASSERT_EQ(request_before.processed_sequence_length, 3); + ASSERT_EQ(cache_before.requests.size(), 1u); + ASSERT_EQ(cache_before.window_blocks.requests.size(), 1u); + ASSERT_EQ(cache_before.window_blocks.blocks_per_request, 2u); + const size_t ring_period = + cache_before.window_blocks.blocks_per_request * + cache_before.block_size; + ASSERT_EQ(ring_period, 4u); + + const size_t reservations_before = + faulting.cache->ReservationCount(); + const size_t commits_before = faulting.cache->CommitCount(); + const size_t releases_before = faulting.cache->ReleaseCount(); + const size_t decodes_before = + faulting.executor->CompletedRealDecodes(); + faulting.executor->FailNextDecode(); + + try { + static_cast(faulting.engine->Step()); + FAIL() << "Expected retryable failure after real windowed decode."; + } catch (const EngineStepError& error) { + EXPECT_EQ(error.Outcome().kind, + StepOutcomeKind::RetryableBatchAbort); + } + + EXPECT_TRUE(faulting.executor->InjectedFailure()); + EXPECT_TRUE(faulting.executor->ObservedRealLogits()); + EXPECT_EQ(faulting.executor->CompletedRealDecodes(), + decodes_before + 1); + EXPECT_EQ(faulting.executor->FailedUnprocessedTokenCount(), 2u); + EXPECT_EQ(faulting.executor->FailedTargetCacheSlots(), 5u); + // The failed run wrote absolute positions [3, 5), which map to ring + // slots [3, 0]. Thus the injected fault occurred only after a real + // continuation run crossed the four-slot ring boundary. + EXPECT_EQ( + static_cast( + request_before.processed_sequence_length) % + ring_period, + ring_period - 1); + EXPECT_EQ( + (faulting.executor->FailedTargetCacheSlots() - 1) % + ring_period, + 0u); + + const auto request_after = request->Snapshot(); + EXPECT_EQ(request_after.status, RequestStatus::Assigned); + EXPECT_EQ(request_after.current_sequence_length, + request_before.current_sequence_length); + EXPECT_EQ(request_after.processed_sequence_length, + request_before.processed_sequence_length); + EXPECT_FALSE(request->HasUnseenTokens()); + EXPECT_TRUE(faulting.engine->HasPendingRequests()); + + const auto cache_after = faulting.cache->Snapshot(); + ExpectCacheOwnershipRestored(cache_after, cache_before); + EXPECT_EQ(faulting.cache->ReservationCount(), + reservations_before + 1); + EXPECT_EQ(faulting.cache->CommitCount(), commits_before); + EXPECT_EQ(faulting.cache->ReleaseCount(), + releases_before + 1); + EXPECT_EQ(faulting.cache->ActiveReservations(), 0u); + EXPECT_EQ(faulting.cache->AbandonedReservationCount(), 0u); + EXPECT_NO_THROW(ThrowIfInvariantsViolated( + cache_after, std::vector{request_after})); + + const auto retry_output = RunTurn(faulting.engine, request); + EXPECT_EQ(retry_output, clean_replay_output); + EXPECT_EQ(retry_output, kSecondTurnOutput); + EXPECT_EQ(std::find(retry_output.begin(), retry_output.end(), + kInvariantFailureToken), + retry_output.end()); + EXPECT_EQ(request->Status(), RequestStatus::TurnComplete); + EXPECT_EQ(request->CurrentSequenceLength(), 14); + EXPECT_EQ(faulting.cache->ActiveReservations(), 0u); + + faulting.engine->RemoveRequest(request); +} + +} // namespace +} // namespace test +} // namespace Generators diff --git a/test/models/engine/synthetic-windowed-multiwrap/decoder.onnx b/test/models/engine/synthetic-windowed-multiwrap/decoder.onnx new file mode 100644 index 0000000000000000000000000000000000000000..e15c8e83c476dfc11f51f31328a64883e7f0278d GIT binary patch literal 10403 zcma(XYjYIG5xwtPpvUNa0wXRiz$y!L5;7RDV+1nc%0@(rPaS_8`2Ih9xDTV= zHYh{*-5}Ep%n8D_e_5Y6a>G*}eu8yx&>M8!G&uL2#76?5=Q~|LJWbCMm`qYPP8%@g zhi%90c45N343cd!(~1UR+l>Q1af00)l&l+gK5Rj5fItRVAAjj4>1Gv1(`a?-zayit z4XQ|-c7x0+Vme;btyg~u6QS4xPzaiYiK7b@{yWH1(NBrn-%b7a zFh2M=aJ#6FGD>DLNGJH_QBj{~kOPq_O(G+l%zE|6Zx1~Gc+lINg~|ut@3(_qvU-gg zIGt)P+t>O;W#HsS)afKXdTQuU z(~lhwTJRE53!QRn+>}&+a8Wgy~i{@;*4J+v@scR{1b1 z{sIQcG_X%zwCRC^8CO+C=&_#z}#w`ORHH}KrjhIg{LAd3Avhd)uGv5(#I^(Xjg3d(GF%u*ilk*5`WKFU39UN=hU<2$c!Z-0PWgSZXkAPF78_ zX3q5&umV8N2AcKra_OG6lbdCFfUOM$(H6n*_Z5G9#cXmg=C-X9g;$w9FRtC@w* zobLomK%ih6ZEA`V>Mr=IxY%gyxY`L~44Rd7()VLyV9UacH(*7O(I1)1lQc*7+fV>x zR$5TdGZkn+cUUrqRF@jZ)J$Kh&CgJRGSr6@$(l9#Qc0jumB(b}ga3(V$quG>Jsy_c zkzvd(nWy6z?gwcQg$l@1$YDYelFU)$iy%FVFlBEAu1G1ipfG^4G&neB*-<+PPgwxZ zf~p1wR|m*-N?G=3`6ly$ySkctWbWMU4*Y_me~^mqkOdB5Q1u5WV1o&o5Kf|~a03iO zsPc|Pk(y*mKbI)#5LptA!x_n;qW&+e;meRjbP?hzvXgi!a#d<*#xB7k41U0P2WujP zumqE>NeV&x>QhHPE5~43w8DlpTt-A-FJ{Z8S8$T5ap~nUVG$8itw2v@`U3q|HG#A@9w)Sj;JMxzmK15plH3r0d~sz<=zs;u zCMz0<%w|NGyhbuPBLNW|Mrqh`%gDT8;gnf%BcY*0V&!F|KGdqEztz^tB#BXIXSS_%oiXgq5OJE-|3-H`pc^vhQ=w2D%{ekz% z>-x0Tn6DV`QSmJh7#wXA+(fitUN?RM`F?pUFXZ6!1O$1tO2U0B zDV>RMMV2HgbL7u?WvZ6-g^Ka^aTdT_FJB&_SD~ftjMJ3}jLLa$>)a7e7gT zg5Udm17_A(g!em9>^ppyR)rdS8>JAUp)3GkWJuOIXTL3vZcY;&HE7q?PX}(?7Q*E~ z>Z6A-{1q(QCs_9wrE0q7$LUy0W#U^^W3iHf>o_bdf-mt#ItF?XTlKv5-r#F3w>$dAB&7O1#; z9QP{~caK!LIgD4qct<}Rg)wjEiG*$k=loKwE+MdoLkLYf-6)EAh2p-1K^~Gf8iQPQd07P(*=w?~U}EAKjokh>m zbo2Zk^4&1*9ULBmK|_3Ik=apbi*FIKPI)-bWL^1l`7z`MVfsAs+?Es5ml$LoWn{>- zP{y7>NJHJz9zC!UQkN!s$~y^O4nBk;5tz%vk_DAhOrLz<(+A`2`2#S)lF4p#8sKTE zKIj0v`y^i`pF(^+kOrRET+vPNaY-6n#BSed;{>D z8~ch&q*=Kw2`^Pf5QJ6YRehCiRq160E#BNHY*tvme1p@h{LX)V+gxM=tW-7Z{eAOE zWgM9nlJMq+G`O^>w6<(FpE68C%G}(LII@*Evfcd8Xl10bfs9qJZvIQcUz32cLr&wk zu$7#|N%1YGPzqfY{UF+D(Ows=oZlJoJuBKd(awvOh;~7=i=tf;?XqZBM7t{5HPO~Y zyDnPkM7PBE7owFuaB@ef(vn3Wl_@d<*VTtxCwt{@f0kjfGD+rOhFovaFwNnnlaeQSsPQ~fDm;ClVL2=AaP_96GD8#`$ zQ>&#CD*p(|KU%5a?3T2TxTmrBw&9!I4#1 zt^gs_Jlya}fUoODD;uqQpqI21Vw#SvuV@~}>V#PPVzIUrRH0d_lS=anUG~S}+T47M m5BNI4C7!H4xr;CRg3#aMVWHpN8t%_qc>Jb8IKZcXQTTsZTTgoc literal 0 HcmV?d00001 diff --git a/test/models/engine/synthetic-windowed-multiwrap/genai_config.json b/test/models/engine/synthetic-windowed-multiwrap/genai_config.json new file mode 100644 index 0000000000..79386a0496 --- /dev/null +++ b/test/models/engine/synthetic-windowed-multiwrap/genai_config.json @@ -0,0 +1,56 @@ +{ + "model": { + "type": "decoder", + "bos_token_id": 0, + "eos_token_id": 1, + "pad_token_id": 0, + "vocab_size": 32, + "context_length": 16, + "decoder": { + "session_options": { + "log_id": "onnxruntime-genai", + "provider_options": [] + }, + "filename": "decoder.onnx", + "num_attention_heads": 1, + "num_key_value_heads": 1, + "head_size": 1, + "hidden_size": 1, + "num_hidden_layers": 2, + "sliding_window": { + "window_size": 3, + "slide_key_value_cache": false, + "slide_inputs": false, + "layers": [ + 1 + ] + }, + "inputs": { + "input_ids": "input_ids", + "block_table": "block_table", + "block_table_windowed": "block_table_windowed", + "cumulative_sequence_lengths": "cumulative_sequence_lengths", + "past_sequence_lengths": "past_sequence_lengths", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value" + }, + "outputs": { + "logits": "logits", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value" + } + } + }, + "search": { + "max_length": 16, + "chunk_size": 2, + "do_sample": false + }, + "engine": { + "dynamic_batching": { + "block_size": 2, + "num_blocks": 8, + "max_batch_size": 1 + } + } +} diff --git a/test/python/create/create_synthetic_windowed_multiwrap_model.py b/test/python/create/create_synthetic_windowed_multiwrap_model.py new file mode 100644 index 0000000000..8fcf95c134 --- /dev/null +++ b/test/python/create/create_synthetic_windowed_multiwrap_model.py @@ -0,0 +1,631 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Create the deterministic sliding-window Engine continuation fixture. + +The graph models two paged KV-cache layers: + +* layer 0 keeps the full sequence; +* layer 1 uses the runtime's repeated sliding-window block table. + +Both layers store token-and-position encodings. Logits read those encodings back +through both block tables, and an invariant-failure token is selected if the +window table does not repeat, its blocks change across continuation, or its +live cache values disagree with the full cache. +""" + +import argparse +import json +import os + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper + +VOCAB_SIZE = 32 +BLOCK_SIZE = 2 +WINDOW_SIZE = 3 +CHUNK_SIZE = 2 +MAX_BATCH_SIZE = 1 +RING_BLOCKS = (CHUNK_SIZE + WINDOW_SIZE - 1 + BLOCK_SIZE - 1) // BLOCK_SIZE +NUM_FULL_BLOCKS = 8 +NUM_WINDOW_BLOCKS = RING_BLOCKS * MAX_BATCH_SIZE +CONTEXT_LENGTH = 16 +EOS_TOKEN_ID = 1 +INVARIANT_FAILURE_TOKEN_ID = 31 + + +def _const(name, array): + tensor = numpy_helper.from_array(np.asarray(array)) + tensor.name = name + return tensor + + +def _decoder_graph(): + def i64(value): + return np.asarray(value, dtype=np.int64) + + full_cache_shape = [NUM_FULL_BLOCKS, BLOCK_SIZE, 1, 1] + window_cache_shape = [NUM_WINDOW_BLOCKS, BLOCK_SIZE, 1, 1] + initializers = [ + _const("c0", i64(0)), + _const("c1", i64(1)), + _const("c2", i64(2)), + _const("c3", i64(3)), + _const("c5", i64(5)), + _const("c7", i64(7)), + _const("c13", i64(13)), + _const("c28", i64(28)), + _const("cB", i64(BLOCK_SIZE)), + _const("cR", i64(RING_BLOCKS)), + _const("cEOS", i64(EOS_TOKEN_ID)), + _const("cInvariantFailure", i64(INVARIANT_FAILURE_TOKEN_ID)), + _const("axis0", i64([0])), + _const("axis1", i64([1])), + _const("start1", i64([1])), + _const("end_all", i64([np.iinfo(np.int64).max])), + _const("flat", i64([-1])), + _const("full_cache_shape", i64(full_cache_shape)), + _const("window_cache_shape", i64(window_cache_shape)), + _const( + "vocab_range", + np.arange(VOCAB_SIZE, dtype=np.int64).reshape(1, VOCAB_SIZE), + ), + ] + + nodes = [] + + def node(op_type, inputs, outputs, *, name=None, **attrs): + if name is not None: + attrs["name"] = name + nodes.append(helper.make_node(op_type, inputs, outputs, **attrs)) + + # Resolve each packed token to a request row and an absolute request position. + node("Shape", ["input_ids"], ["ids_shape"]) + node("Squeeze", ["ids_shape"], ["num_tokens"]) + node("Range", ["c0", "num_tokens", "c1"], ["token_index"]) + node( + "Slice", + ["cumulative_sequence_lengths", "start1", "end_all", "axis0"], + ["boundaries_i32"], + ) + node("Cast", ["boundaries_i32"], ["boundaries"], to=TensorProto.INT64) + node("Unsqueeze", ["token_index", "axis1"], ["token_index_col"]) + node("Unsqueeze", ["boundaries", "axis0"], ["boundaries_row"]) + node("GreaterOrEqual", ["token_index_col", "boundaries_row"], ["at_or_past"]) + node("Cast", ["at_or_past"], ["at_or_past_i64"], to=TensorProto.INT64) + node("ReduceSum", ["at_or_past_i64", "axis1"], ["row_id"], keepdims=0) + + node( + "Cast", + ["cumulative_sequence_lengths"], + ["cumulative_sequence_lengths_i64"], + to=TensorProto.INT64, + ) + node( + "Gather", + ["cumulative_sequence_lengths_i64", "row_id"], + ["row_start"], + axis=0, + ) + node("Sub", ["token_index", "row_start"], ["offset_in_row"]) + node( + "Cast", + ["past_sequence_lengths"], + ["past_sequence_lengths_i64"], + to=TensorProto.INT64, + ) + node( + "Gather", + ["past_sequence_lengths_i64", "row_id"], + ["past_of_row"], + axis=0, + ) + node("Add", ["past_of_row", "offset_in_row"], ["pos"]) + node("Sub", ["pos", "c1"], ["previous_pos_unclamped"]) + node("Max", ["previous_pos_unclamped", "c0"], ["previous_pos"]) + + node("Cast", ["block_table"], ["block_table_i64"], to=TensorProto.INT64) + node( + "Cast", + ["block_table_windowed"], + ["block_table_windowed_i64"], + to=TensorProto.INT64, + ) + node("Unsqueeze", ["row_id", "axis1"], ["row_id_col"]) + + def map_positions(prefix, positions, table): + block_col = f"{prefix}_block_col" + block_col_base = f"{prefix}_block_col_base" + slot = f"{prefix}_slot_in_block" + block_col_col = f"{prefix}_block_col_col" + gather_index = f"{prefix}_block_gather_index" + block_id = f"{prefix}_block_id" + block_base = f"{prefix}_block_base" + physical = f"{prefix}_physical" + + node("Div", [positions, "cB"], [block_col]) + node("Mul", [block_col, "cB"], [block_col_base]) + node("Sub", [positions, block_col_base], [slot]) + node("Unsqueeze", [block_col, "axis1"], [block_col_col]) + node("Concat", ["row_id_col", block_col_col], [gather_index], axis=1) + node("GatherND", [table, gather_index], [block_id]) + node("Mul", [block_id, "cB"], [block_base]) + node("Add", [block_base, slot], [physical]) + return block_col, block_id, physical + + current_block_col, _, current_full_physical = map_positions("current_full", "pos", "block_table_i64") + _, current_window_block, current_window_physical = map_positions( + "current_window", "pos", "block_table_windowed_i64" + ) + _, _, previous_full_physical = map_positions("previous_full", "previous_pos", "block_table_i64") + _, _, previous_window_physical = map_positions("previous_window", "previous_pos", "block_table_windowed_i64") + + # Position zero remains in the full cache. Its value records the first + # window block id so continuation can prove that the ring stayed resident. + node("Gather", ["block_table_i64", "c0"], ["first_full_block_per_row"], axis=1) + node( + "Gather", + ["first_full_block_per_row", "row_id"], + ["first_full_block"], + axis=0, + ) + node("Mul", ["first_full_block", "cB"], ["first_full_physical"]) + node( + "Gather", + ["block_table_windowed_i64", "c0"], + ["first_window_block_per_row"], + axis=1, + ) + node( + "Gather", + ["first_window_block_per_row", "row_id"], + ["first_window_block"], + axis=0, + ) + + # Store exact integer-valued float encodings. Including the absolute + # position makes a continuation that resets past_sequence_lengths diverge. + node("Mul", ["input_ids", "c7"], ["key_token_term"]) + node("Mul", ["pos", "c3"], ["key_position_term"]) + node("Add", ["key_token_term", "key_position_term"], ["key_without_bias"]) + node("Add", ["key_without_bias", "c1"], ["key_encoding_i64"]) + node("Cast", ["key_encoding_i64"], ["key_encoding"], to=TensorProto.FLOAT) + + node("Mul", ["input_ids", "c5"], ["value_token_term"]) + node("Mul", ["pos", "c2"], ["value_position_term"]) + node( + "Add", + ["value_token_term", "value_position_term"], + ["value_without_bias"], + ) + node("Add", ["value_without_bias", "c2"], ["value_encoding_i64"]) + node( + "Cast", + ["value_encoding_i64"], + ["window_value_encoding"], + to=TensorProto.FLOAT, + ) + node( + "Cast", + [current_window_block], + ["window_owner_encoding"], + to=TensorProto.FLOAT, + ) + + node( + "Reshape", + ["past_key_values.0.key", "flat"], + ["past_full_key_flat"], + ) + node( + "Reshape", + ["past_key_values.0.value", "flat"], + ["past_full_value_flat"], + ) + node( + "Reshape", + ["past_key_values.1.key", "flat"], + ["past_window_key_flat"], + ) + node( + "Reshape", + ["past_key_values.1.value", "flat"], + ["past_window_value_flat"], + ) + node( + "Unsqueeze", + [current_full_physical, "axis1"], + ["current_full_scatter_index"], + ) + node( + "Unsqueeze", + [current_window_physical, "axis1"], + ["current_window_scatter_index"], + ) + + node( + "ScatterND", + ["past_full_key_flat", "current_full_scatter_index", "key_encoding"], + ["present_full_key_flat"], + ) + node( + "ScatterND", + [ + "past_full_value_flat", + "current_full_scatter_index", + "window_owner_encoding", + ], + ["present_full_value_flat"], + ) + node( + "ScatterND", + ["past_window_key_flat", "current_window_scatter_index", "key_encoding"], + ["present_window_key_flat"], + ) + node( + "ScatterND", + [ + "past_window_value_flat", + "current_window_scatter_index", + "window_value_encoding", + ], + ["present_window_value_flat"], + ) + + node( + "Reshape", + ["present_full_key_flat", "full_cache_shape"], + ["present.0.key"], + ) + node( + "Reshape", + ["present_full_value_flat", "full_cache_shape"], + ["present.0.value"], + ) + node( + "Reshape", + ["present_window_key_flat", "window_cache_shape"], + ["present.1.key"], + ) + node( + "Reshape", + ["present_window_value_flat", "window_cache_shape"], + ["present.1.value"], + ) + + # Logits consume values read through both cache layers, including the + # previous live window position after the ring has wrapped. + node( + "Gather", + ["present_full_key_flat", "first_full_physical"], + ["read_full_first_key"], + axis=0, + name="read_full_first_key", + ) + node( + "Gather", + ["present_full_key_flat", current_full_physical], + ["read_full_current_key"], + axis=0, + name="read_full_current_key", + ) + node( + "Gather", + ["present_full_key_flat", previous_full_physical], + ["read_full_previous_key"], + axis=0, + name="read_full_previous_key", + ) + node( + "Gather", + ["present_full_value_flat", "first_full_physical"], + ["read_first_window_owner"], + axis=0, + name="read_first_window_owner", + ) + node( + "Gather", + ["present_window_key_flat", previous_window_physical], + ["read_window_previous_key"], + axis=0, + name="read_window_previous_key", + ) + node( + "Gather", + ["present_window_key_flat", current_window_physical], + ["read_window_current_key"], + axis=0, + name="read_window_current_key", + ) + node( + "Gather", + ["present_window_value_flat", previous_window_physical], + ["read_window_previous_value"], + axis=0, + name="read_window_previous_value", + ) + node( + "Gather", + ["present_window_value_flat", current_window_physical], + ["read_window_current_value"], + axis=0, + name="read_window_current_value", + ) + + # The current and prior-cycle columns must name the same physical block. + node("Sub", [current_block_col, "cR"], ["prior_cycle_col_unclamped"]) + node("Max", ["prior_cycle_col_unclamped", "c0"], ["prior_cycle_col"]) + node("Unsqueeze", ["prior_cycle_col", "axis1"], ["prior_cycle_col_col"]) + node( + "Concat", + ["row_id_col", "prior_cycle_col_col"], + ["prior_cycle_gather_index"], + axis=1, + ) + node( + "GatherND", + ["block_table_windowed_i64", "prior_cycle_gather_index"], + ["prior_cycle_window_block"], + ) + node( + "GreaterOrEqual", + [current_block_col, "cR"], + ["has_prior_block_cycle"], + ) + node( + "Equal", + [current_window_block, "prior_cycle_window_block"], + ["window_block_repeats"], + ) + node("Not", ["has_prior_block_cycle"], ["before_first_block_cycle"]) + node( + "Or", + ["before_first_block_cycle", "window_block_repeats"], + ["repeated_window_block_valid"], + name="guard_repeated_window_block", + ) + + node( + "Equal", + ["read_full_previous_key", "read_window_previous_key"], + ["previous_cache_values_match"], + ) + node( + "Equal", + ["read_full_current_key", "read_window_current_key"], + ["current_cache_values_match"], + ) + node( + "Cast", + ["first_window_block"], + ["first_window_block_f"], + to=TensorProto.FLOAT, + ) + node( + "Equal", + ["read_first_window_owner", "first_window_block_f"], + ["window_owner_stable"], + name="guard_stable_window_owner", + ) + node( + "And", + ["previous_cache_values_match", "current_cache_values_match"], + ["cache_values_match"], + ) + node( + "And", + ["cache_values_match", "window_owner_stable"], + ["cache_and_owner_valid"], + ) + node( + "And", + ["cache_and_owner_valid", "repeated_window_block_valid"], + ["window_invariants_valid"], + ) + + score_terms = [ + "read_full_first_key", + "read_full_current_key", + "read_window_previous_key", + "read_window_current_key", + "read_window_previous_value", + "read_window_current_value", + ] + score = score_terms[0] + for index, term in enumerate(score_terms[1:], start=1): + output = f"score_sum_{index}" + node("Add", [score, term], [output]) + score = output + node("Cast", [score], ["score_i64"], to=TensorProto.INT64) + node("Div", ["score_i64", "c28"], ["score_div"]) + node("Mul", ["score_div", "c28"], ["score_floor"]) + node("Sub", ["score_i64", "score_floor"], ["score_mod"]) + node("Add", ["score_mod", "c2"], ["normal_next_token"]) + node( + "Where", + ["window_invariants_valid", "normal_next_token", "cInvariantFailure"], + ["guarded_next_token"], + ) + + # EOS at absolute positions 2 and 13 creates two short turns while leaving + # max_length headroom. The continuation spans positions 3..11; normal + # outputs at 11 and 12 validate both columns of the two-block ring. + node("Equal", ["pos", "c2"], ["is_first_turn_eos_position"]) + node("Equal", ["pos", "c13"], ["is_second_turn_eos_position"]) + node( + "Or", + ["is_first_turn_eos_position", "is_second_turn_eos_position"], + ["is_eos_position"], + ) + node( + "Where", + ["is_eos_position", "cEOS", "guarded_next_token"], + ["next_token"], + ) + + node("Unsqueeze", ["next_token", "axis1"], ["next_token_col"]) + node("Equal", ["next_token_col", "vocab_range"], ["is_next_per_token"]) + node("Sub", ["boundaries", "c1"], ["last_token_index"]) + node( + "Gather", + ["is_next_per_token", "last_token_index"], + ["is_next_per_request"], + axis=0, + ) + node( + "Cast", + ["is_next_per_request"], + ["logits"], + to=TensorProto.FLOAT16, + ) + + inputs = [ + helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["num_tokens"]), + helper.make_tensor_value_info( + "cumulative_sequence_lengths", + TensorProto.INT32, + ["batch_plus_1"], + ), + helper.make_tensor_value_info("past_sequence_lengths", TensorProto.INT32, ["batch"]), + helper.make_tensor_value_info("block_table", TensorProto.INT32, ["batch", "max_blocks"]), + helper.make_tensor_value_info( + "block_table_windowed", + TensorProto.INT32, + ["batch", "max_blocks"], + ), + helper.make_tensor_value_info( + "past_key_values.0.key", + TensorProto.FLOAT, + full_cache_shape, + ), + helper.make_tensor_value_info( + "past_key_values.0.value", + TensorProto.FLOAT, + full_cache_shape, + ), + helper.make_tensor_value_info( + "past_key_values.1.key", + TensorProto.FLOAT, + window_cache_shape, + ), + helper.make_tensor_value_info( + "past_key_values.1.value", + TensorProto.FLOAT, + window_cache_shape, + ), + ] + outputs = [ + helper.make_tensor_value_info("logits", TensorProto.FLOAT16, ["batch_size", VOCAB_SIZE]), + helper.make_tensor_value_info("present.0.key", TensorProto.FLOAT, full_cache_shape), + helper.make_tensor_value_info("present.0.value", TensorProto.FLOAT, full_cache_shape), + helper.make_tensor_value_info("present.1.key", TensorProto.FLOAT, window_cache_shape), + helper.make_tensor_value_info("present.1.value", TensorProto.FLOAT, window_cache_shape), + ] + return helper.make_graph( + nodes, + "synthetic_windowed_multiwrap_decoder", + inputs, + outputs, + initializer=initializers, + ) + + +def create_decoder(output_dir): + model = helper.make_model( + _decoder_graph(), + opset_imports=[helper.make_operatorsetid("", 17)], + ir_version=9, + producer_name="onnxruntime-genai", + producer_version="0.0.0", + ) + metadata = model.metadata_props.add() + metadata.key = "fixture" + metadata.value = "engine-windowed-multiwrap-continuation" + onnx.checker.check_model(model) + onnx.save_model(model, os.path.join(output_dir, "decoder.onnx")) + + +def create_config(output_dir): + config = { + "model": { + "type": "decoder", + "bos_token_id": 0, + "eos_token_id": EOS_TOKEN_ID, + "pad_token_id": 0, + "vocab_size": VOCAB_SIZE, + "context_length": CONTEXT_LENGTH, + "decoder": { + "session_options": { + "log_id": "onnxruntime-genai", + "provider_options": [], + }, + "filename": "decoder.onnx", + "num_attention_heads": 1, + "num_key_value_heads": 1, + "head_size": 1, + "hidden_size": 1, + "num_hidden_layers": 2, + "sliding_window": { + "window_size": WINDOW_SIZE, + "slide_key_value_cache": False, + "slide_inputs": False, + "layers": [1], + }, + "inputs": { + "input_ids": "input_ids", + "block_table": "block_table", + "block_table_windowed": "block_table_windowed", + "cumulative_sequence_lengths": "cumulative_sequence_lengths", + "past_sequence_lengths": "past_sequence_lengths", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value", + }, + "outputs": { + "logits": "logits", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value", + }, + }, + }, + "search": { + "max_length": CONTEXT_LENGTH, + "chunk_size": CHUNK_SIZE, + "do_sample": False, + }, + "engine": { + "dynamic_batching": { + "block_size": BLOCK_SIZE, + "num_blocks": NUM_FULL_BLOCKS, + "max_batch_size": MAX_BATCH_SIZE, + }, + }, + } + with open(os.path.join(output_dir, "genai_config.json"), "w") as config_file: + json.dump(config, config_file, indent=2) + config_file.write("\n") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--output_dir", + default=os.path.join( + os.path.dirname(__file__), + "..", + "..", + "models", + "engine", + "synthetic-windowed-multiwrap", + ), + ) + args = parser.parse_args() + output_dir = os.path.normpath(args.output_dir) + os.makedirs(output_dir, exist_ok=True) + create_decoder(output_dir) + create_config(output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/python/models/test_engine_windowed_multiwrap.py b/test/python/models/test_engine_windowed_multiwrap.py new file mode 100644 index 0000000000..454f01ac1b --- /dev/null +++ b/test/python/models/test_engine_windowed_multiwrap.py @@ -0,0 +1,199 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +"""Executable Engine coverage for continuation across window-ring wraps.""" + +import json +from itertools import pairwise +from pathlib import Path + +import numpy as np +import onnx +import onnxruntime_genai as og +import pytest + +_MODEL_SUBPATH = Path("engine") / "synthetic-windowed-multiwrap" + +_VOCAB_SIZE = 32 +_BLOCK_SIZE = 2 +_WINDOW_SIZE = 3 +_CHUNK_SIZE = 2 +_RING_BLOCKS = 2 +_RING_PERIOD = _RING_BLOCKS * _BLOCK_SIZE +_NUM_FULL_BLOCKS = 8 +_MAX_LENGTH = 16 +_EOS_TOKEN_ID = 1 +_INVARIANT_FAILURE_TOKEN_ID = 31 + +_INITIAL_PROMPT = [4, 6] +_EXPECTED_FIRST_TURN = [12] +_CONTINUATION = [8, 3, 10, 5, 12, 7, 14, 9, 16] +_EXPECTED_SECOND_TURN = [28, 17] + +_DEVICES = ["cpu"] + (["cuda"] if og.is_cuda_available() else []) + + +def _fixture_path(test_data_path) -> Path: + if not test_data_path: + pytest.skip("--test_models is required for the synthetic Engine fixture") + path = Path(test_data_path) / _MODEL_SUBPATH + if not path.exists(): + pytest.fail(f"synthetic Engine fixture is missing: {path}") + return path + + +def _cache_key(token, position): + return token * 7 + position * 3 + 1 + + +def _cache_value(token, position): + return token * 5 + position * 2 + 2 + + +def _normal_token(sequence, position): + """Mirror the graph's cache-read score for a non-EOS position.""" + first_key = _cache_key(sequence[0], 0) + previous_key = _cache_key(sequence[position - 1], position - 1) + current_key = _cache_key(sequence[position], position) + previous_value = _cache_value(sequence[position - 1], position - 1) + current_value = _cache_value(sequence[position], position) + score = first_key + current_key + previous_key + current_key + previous_value + current_value + return score % 28 + 2 + + +def _wrap_count(start_position, token_count): + positions = range(start_position, start_position + token_count) + slots = [position % _RING_PERIOD for position in positions] + return sum(current < previous for previous, current in pairwise(slots)) + + +def _new_request(engine, model, prompt): + params = og.GeneratorParams(model) + params.set_search_options(do_sample=False, max_length=_MAX_LENGTH) + request = og.Request(params) + request.add_tokens(np.asarray(prompt, dtype=np.int32)) + engine.add_request(request) + return request + + +def _run_turn(engine, request): + tokens = [] + steps = 0 + while not request.is_turn_complete(): + ready = engine.step() + steps += 1 + assert steps < 100, "synthetic turn did not reach its absolute-position EOS" + if ready is None: + # Partial prefill chunks commit cache progress without sampling. + continue + assert ready is request + while ready.has_unseen_tokens(): + tokens.append(ready.get_unseen_token()) + return tokens + + +def test_windowed_multiwrap_fixture_schema(test_data_path): + model_path = _fixture_path(test_data_path) + config = json.loads((model_path / "genai_config.json").read_text(encoding="utf-8")) + decoder = config["model"]["decoder"] + dynamic = config["engine"]["dynamic_batching"] + + assert decoder["num_hidden_layers"] == 2 + assert decoder["sliding_window"] == { + "window_size": _WINDOW_SIZE, + "slide_key_value_cache": False, + "slide_inputs": False, + "layers": [1], + } + assert decoder["inputs"]["block_table_windowed"] == "block_table_windowed" + assert config["model"]["vocab_size"] == _VOCAB_SIZE + assert config["model"]["eos_token_id"] == _EOS_TOKEN_ID + assert config["search"]["max_length"] == _MAX_LENGTH + assert config["search"]["chunk_size"] == _CHUNK_SIZE + assert dynamic["block_size"] == _BLOCK_SIZE + assert dynamic["num_blocks"] == _NUM_FULL_BLOCKS + assert dynamic["max_batch_size"] == 1 + + assert (_CHUNK_SIZE + _WINDOW_SIZE - 1 + _BLOCK_SIZE - 1) // _BLOCK_SIZE == _RING_BLOCKS + continuation_start = len(_INITIAL_PROMPT) + len(_EXPECTED_FIRST_TURN) + assert len(_CONTINUATION) > 2 * _RING_PERIOD + assert _wrap_count(continuation_start, len(_CONTINUATION)) == 2 + + graph_model = onnx.load(model_path / "decoder.onnx", load_external_data=False) + onnx.checker.check_model(graph_model) + metadata = {entry.key: entry.value for entry in graph_model.metadata_props} + assert metadata["fixture"] == "engine-windowed-multiwrap-continuation" + + input_shapes = { + value.name: [dimension.dim_value for dimension in value.type.tensor_type.shape.dim] + for value in graph_model.graph.input + if value.name.startswith("past_key_values") + } + assert input_shapes["past_key_values.0.key"] == [_NUM_FULL_BLOCKS, 2, 1, 1] + assert input_shapes["past_key_values.0.value"] == [_NUM_FULL_BLOCKS, 2, 1, 1] + assert input_shapes["past_key_values.1.key"] == [2, 2, 1, 1] + assert input_shapes["past_key_values.1.value"] == [2, 2, 1, 1] + + node_names = {node.name for node in graph_model.graph.node} + assert { + "read_full_current_key", + "read_full_previous_key", + "read_window_previous_key", + "read_window_current_value", + "guard_repeated_window_block", + "guard_stable_window_owner", + } <= node_names + + +@pytest.mark.parametrize("device", _DEVICES) +def test_continuation_crosses_two_window_ring_wraps_and_matches_clean_replay(test_data_path, device): + model_path = _fixture_path(test_data_path) + config = og.Config(str(model_path)) + config.clear_providers() + if device == "cuda": + config.append_provider("cuda") + model = og.Model(config) + + # The first normal result is followed by EOS at absolute position 2, well + # below the session max. EOS is not part of the retained logical sequence. + assert _normal_token(_INITIAL_PROMPT, 1) == _EXPECTED_FIRST_TURN[0] + engine = og.Engine(model) + request = _new_request(engine, model, _INITIAL_PROMPT) + first_turn = _run_turn(engine, request) + + assert first_turn == _EXPECTED_FIRST_TURN + assert request.status == og.RequestStatus.TURN_COMPLETE + first_turn_length = len(_INITIAL_PROMPT) + len(first_turn) + assert first_turn_length == 3 + assert first_turn_length < _MAX_LENGTH + + # Positions 3..11 traverse the four-slot ring twice. The graph returns 31 + # if the repeated table, retained ring ownership, or values read through + # the full/windowed caches disagree. Tokens at positions 11 and 12 check + # both columns of the two-block ring before EOS at position 13. + request.continue_with(np.asarray(_CONTINUATION, dtype=np.int32)) + assert request.status == og.RequestStatus.QUEUED + second_turn = _run_turn(engine, request) + + assert second_turn == _EXPECTED_SECOND_TURN + assert _INVARIANT_FAILURE_TOKEN_ID not in second_turn + assert request.status == og.RequestStatus.TURN_COMPLETE + retained_length = first_turn_length + len(_CONTINUATION) + len(second_turn) + assert retained_length == 14 + assert retained_length < _MAX_LENGTH + engine.remove_request(request) + + # Recompute the same logical context in a fresh cache. Exact parity proves + # that continuation's retained full cache and twice-wrapped window cache + # agree with a clean chunked replay at absolute positions 0..13. + replay_prompt = _INITIAL_PROMPT + first_turn + _CONTINUATION + assert _normal_token(replay_prompt, len(replay_prompt) - 1) == _EXPECTED_SECOND_TURN[0] + replay_after_first_output = replay_prompt + _EXPECTED_SECOND_TURN[:1] + assert _normal_token(replay_after_first_output, len(replay_after_first_output) - 1) == _EXPECTED_SECOND_TURN[1] + replay_engine = og.Engine(model) + replay_request = _new_request(replay_engine, model, replay_prompt) + replay_turn = _run_turn(replay_engine, replay_request) + + assert replay_turn == second_turn == _EXPECTED_SECOND_TURN + assert _INVARIANT_FAILURE_TOKEN_ID not in replay_turn + replay_engine.remove_request(replay_request) From d4d0b41edb1dc85f330f9c9e057abf4021839f12 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 19:14:14 -0500 Subject: [PATCH 07/19] Test Python orphaned request reclamation Fill every retained Engine slot, release all public request handles without explicit removal, and verify the next admission reclaims capacity and completes. File: test/python/test_onnxruntime_genai_engine.py Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b232f82c-f25c-429a-a855-7f8c8f40bbf3 --- test/python/test_onnxruntime_genai_engine.py | 27 ++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/test/python/test_onnxruntime_genai_engine.py b/test/python/test_onnxruntime_genai_engine.py index a8b192bc14..035cc03b1a 100644 --- a/test/python/test_onnxruntime_genai_engine.py +++ b/test/python/test_onnxruntime_genai_engine.py @@ -345,6 +345,33 @@ def test_request_lifecycle_status(model): assert request.status == og.RequestStatus.CLOSED +def test_last_handle_release_reclaims_retained_capacity(model): + engine = og.Engine(model) + sinks = [_Sink() for _ in range(8)] + requests = [ + _add_request(engine, model, [5 + index, 9, 13], 1, sinks[index]) + for index in range(8) + ] + + while not all(request.is_turn_complete() for request in requests): + ready = engine.step() + assert ready is not None + _drain(ready) + + # Every TurnComplete request still owns one of the eight resident slots. Dropping all public + # handles must mark them abandoned so the next admission can reclaim that capacity. + requests.clear() + del ready + gc.collect() + + replacement_sink = _Sink() + replacement = _add_request(engine, model, _PROMPT_A, 4, replacement_sink) + _run(engine) + + assert replacement_sink.tokens == predicted_tokens(_PROMPT_A, 4) + assert replacement.status == og.RequestStatus.CLOSED + + def test_remove_request_freezes_output(model): max_new = 40 sibling_new = 16 From f044a25230383a4f9499e77a08e72264d9034354 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 20:29:09 -0500 Subject: [PATCH 08/19] Retain request handles through Engine execution Update synthetic and real-model Engine tests to keep public Request handles alive until completion, matching automatic orphan cancellation semantics. Use precise is_turn_complete checks and correct static cache lifecycle terminology. Files: src/engine/cache_manager.cpp; test/python/test_onnxruntime_genai_engine.py; test/python/integration/test_integration_engine.py; examples/python/engine/continuous-batching.py Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b232f82c-f25c-429a-a855-7f8c8f40bbf3 --- examples/python/engine/continuous-batching.py | 2 +- src/engine/cache_manager.cpp | 2 +- .../integration/test_integration_engine.py | 38 +++++++++++++------ test/python/test_onnxruntime_genai_engine.py | 31 ++++++++++----- 4 files changed, 49 insertions(+), 24 deletions(-) diff --git a/examples/python/engine/continuous-batching.py b/examples/python/engine/continuous-batching.py index 263fb2f924..bf79ad9794 100644 --- a/examples/python/engine/continuous-batching.py +++ b/examples/python/engine/continuous-batching.py @@ -84,7 +84,7 @@ def drain(self, request: og.Request): token = request.get_unseen_token() client_request.token_stream += client_request.streaming_tokenizer.decode(token) - if request.is_done(): + if request.is_turn_complete(): assert client_request is not None, "Client request not found in the pool" if self.debug: diff --git a/src/engine/cache_manager.cpp b/src/engine/cache_manager.cpp index bf240d2243..00b82b207f 100644 --- a/src/engine/cache_manager.cpp +++ b/src/engine/cache_manager.cpp @@ -101,7 +101,7 @@ void StaticCacheManager::Allocate(const std::vector>& r return IsTurnComplete(request->status_) || IsClosed(request->status_); })) { - // If all requests are completed, we can deallocate them before allocating the new requests. + // If every request is TurnComplete or Closed, recycle the static batch before allocating new requests. Deallocate(cache_allocated_requests_); } diff --git a/test/python/integration/test_integration_engine.py b/test/python/integration/test_integration_engine.py index 68fbe90eab..365bfe20b3 100644 --- a/test/python/integration/test_integration_engine.py +++ b/test/python/integration/test_integration_engine.py @@ -94,7 +94,7 @@ def _drain(ready) -> bool: sink = ready.get_opaque_data() while ready.has_unseen_tokens(): sink.tokens.append(ready.get_unseen_token()) - return ready.is_done() + return ready.is_turn_complete() def _run(engine, *, max_steps=_MAX_STEPS) -> None: @@ -111,8 +111,9 @@ def _run(engine, *, max_steps=_MAX_STEPS) -> None: def _generate_isolated(model, prompt_tokens, max_new_tokens, *, min_new_tokens=0) -> list[int]: sink = _Sink() engine = og.Engine(model) - _add_request(engine, model, prompt_tokens, max_new_tokens, sink, min_new_tokens=min_new_tokens) + request = _add_request(engine, model, prompt_tokens, max_new_tokens, sink, min_new_tokens=min_new_tokens) _run(engine) + assert request.status == og.RequestStatus.CLOSED del engine gc.collect() return sink.tokens @@ -192,7 +193,7 @@ def test_staggered_admission(bundle): engine = og.Engine(bundle.model) sink_a = _Sink() - _add_request(engine, bundle.model, prompt_a, max_new, sink_a) + request_a = _add_request(engine, bundle.model, prompt_a, max_new, sink_a) for _ in range(3): if not engine.has_pending_requests(): @@ -205,12 +206,14 @@ def test_staggered_admission(bundle): assert len(sink_a.tokens) > 0, "first request produced nothing before staggered admission" sink_b = _Sink() - _add_request(engine, bundle.model, prompt_b, max_new, sink_b) + request_b = _add_request(engine, bundle.model, prompt_b, max_new, sink_b) _run(engine) assert sink_a.tokens == isolated_a assert sink_b.tokens == isolated_b + assert request_a.status == og.RequestStatus.CLOSED + assert request_b.status == og.RequestStatus.CLOSED def test_isolated_matches_batched(bundle): @@ -222,11 +225,14 @@ def test_isolated_matches_batched(bundle): engine = og.Engine(bundle.model) sinks = {p: _Sink() for p in _PROMPTS} - for p, sink in sinks.items(): + requests = [ _add_request(engine, bundle.model, bundle.tokenizer.encode(p), max_new, sink) + for p, sink in sinks.items() + ] _run(engine) assert sinks[prompt].tokens == isolated, "batched output diverged from the isolated run" + assert all(request.status == og.RequestStatus.CLOSED for request in requests) def test_output_isolation(bundle): @@ -239,12 +245,15 @@ def test_output_isolation(bundle): engine = og.Engine(bundle.model) s0, s1 = _Sink(), _Sink() - _add_request(engine, bundle.model, bundle.tokenizer.encode(p0), max_new, s0) - _add_request(engine, bundle.model, bundle.tokenizer.encode(p1), max_new, s1) + requests = [ + _add_request(engine, bundle.model, bundle.tokenizer.encode(p0), max_new, s0), + _add_request(engine, bundle.model, bundle.tokenizer.encode(p1), max_new, s1), + ] _run(engine) assert s0.tokens == isolated0 assert s1.tokens == isolated1 + assert all(request.status == og.RequestStatus.CLOSED for request in requests) def test_completion_isolation(bundle): @@ -255,15 +264,17 @@ def test_completion_isolation(bundle): engine = og.Engine(bundle.model) short_sink, long_sink = _Sink(), _Sink() - _add_request( + short_request = _add_request( engine, bundle.model, bundle.tokenizer.encode(short_prompt), short_new, short_sink, min_new_tokens=short_new, ) - _add_request(engine, bundle.model, bundle.tokenizer.encode(long_prompt), long_new, long_sink) + long_request = _add_request(engine, bundle.model, bundle.tokenizer.encode(long_prompt), long_new, long_sink) _run(engine) assert len(short_sink.tokens) == short_new, "forced-length request did not stop at its bound" assert long_sink.tokens == long_isolated, "survivor diverged after its sibling completed" + assert short_request.status == og.RequestStatus.CLOSED + assert long_request.status == og.RequestStatus.CLOSED def test_max_length_stops(bundle): @@ -309,7 +320,7 @@ def test_remove_request_stops_output(bundle): sink_a, sink_b = _Sink(), _Sink() request_a = _add_request(engine, bundle.model, bundle.tokenizer.encode(_PROMPTS[0]), max_new, sink_a) - _add_request(engine, bundle.model, sibling_prompt, max_new, sink_b) + request_b = _add_request(engine, bundle.model, sibling_prompt, max_new, sink_b) for _ in range(4): if not engine.has_pending_requests(): @@ -328,6 +339,7 @@ def test_remove_request_stops_output(bundle): assert len(sink_a.tokens) == frozen_a, "removed request kept producing tokens" assert sink_b.tokens == sibling_isolated, "sibling diverged after request removal" + assert request_b.status == og.RequestStatus.CLOSED def test_engine_teardown_and_recreation(bundle): @@ -337,15 +349,17 @@ def test_engine_teardown_and_recreation(bundle): first = og.Engine(bundle.model) sink1 = _Sink() - _add_request(first, bundle.model, prompt_tokens, max_new, sink1) + first_request = _add_request(first, bundle.model, prompt_tokens, max_new, sink1) _run(first) assert sink1.tokens == expected + assert first_request.status == og.RequestStatus.CLOSED del first gc.collect() second = og.Engine(bundle.model) assert not second.has_pending_requests() sink2 = _Sink() - _add_request(second, bundle.model, prompt_tokens, max_new, sink2) + second_request = _add_request(second, bundle.model, prompt_tokens, max_new, sink2) _run(second) assert sink2.tokens == expected + assert second_request.status == og.RequestStatus.CLOSED diff --git a/test/python/test_onnxruntime_genai_engine.py b/test/python/test_onnxruntime_genai_engine.py index 035cc03b1a..dd0b554729 100644 --- a/test/python/test_onnxruntime_genai_engine.py +++ b/test/python/test_onnxruntime_genai_engine.py @@ -116,8 +116,9 @@ def _run(engine, *, max_steps=_MAX_STEPS): def _generate_isolated(model, prompt, max_new_tokens): sink = _Sink() engine = og.Engine(model) - _add_request(engine, model, prompt, max_new_tokens, sink) + request = _add_request(engine, model, prompt, max_new_tokens, sink) _run(engine) + assert request.status == og.RequestStatus.CLOSED del engine gc.collect() return sink.tokens @@ -166,13 +167,16 @@ def test_isolated_matches_simultaneous(model): engine = og.Engine(model) sink_a, sink_b = _Sink(), _Sink() - _add_request(engine, model, _PROMPT_A, max_new, sink_a) - _add_request(engine, model, _PROMPT_B, max_new, sink_b) + requests = [ + _add_request(engine, model, _PROMPT_A, max_new, sink_a), + _add_request(engine, model, _PROMPT_B, max_new, sink_b), + ] assert engine.has_pending_requests() _run(engine) assert sink_a.tokens == isolated_a, "request A diverged when batched with B" assert sink_b.tokens == isolated_b, "request B diverged when batched with A" + assert all(request.status == og.RequestStatus.CLOSED for request in requests) def test_staggered_admission(model): @@ -182,7 +186,7 @@ def test_staggered_admission(model): engine = og.Engine(model) sink_a = _Sink() - _add_request(engine, model, _PROMPT_A, max_new, sink_a) + request_a = _add_request(engine, model, _PROMPT_A, max_new, sink_a) for _ in range(3): if not engine.has_pending_requests(): @@ -191,11 +195,13 @@ def test_staggered_admission(model): assert len(sink_a.tokens) > 0, "first request produced nothing before staggered admission" sink_b = _Sink() - _add_request(engine, model, _PROMPT_B, max_new, sink_b) + request_b = _add_request(engine, model, _PROMPT_B, max_new, sink_b) _run(engine) assert sink_a.tokens == expected_a assert sink_b.tokens == expected_b + assert request_a.status == og.RequestStatus.CLOSED + assert request_b.status == og.RequestStatus.CLOSED def test_max_length_stops(model): @@ -228,12 +234,14 @@ def test_completion_isolation(model): engine = og.Engine(model) short_sink, long_sink = _Sink(), _Sink() - _add_request(engine, model, _PROMPT_A, short_new, short_sink) - _add_request(engine, model, _PROMPT_LONG, long_new, long_sink) + short_request = _add_request(engine, model, _PROMPT_A, short_new, short_sink) + long_request = _add_request(engine, model, _PROMPT_LONG, long_new, long_sink) _run(engine) assert short_sink.tokens == predicted_tokens(_PROMPT_A, short_new) assert long_sink.tokens == long_isolated, "survivor diverged after its sibling completed" + assert short_request.status == og.RequestStatus.CLOSED + assert long_request.status == og.RequestStatus.CLOSED def test_continuation_while_peer_remains_active(model): @@ -380,7 +388,7 @@ def test_remove_request_freezes_output(model): engine = og.Engine(model) sink_a, sink_b = _Sink(), _Sink() request_a = _add_request(engine, model, _PROMPT_A, max_new, sink_a) - _add_request(engine, model, _PROMPT_B, sibling_new, sink_b) + request_b = _add_request(engine, model, _PROMPT_B, sibling_new, sink_b) for _ in range(4): if not engine.has_pending_requests(): @@ -395,6 +403,7 @@ def test_remove_request_freezes_output(model): assert sink_a.tokens == frozen_a, "removed request kept producing tokens" assert sink_b.tokens == sibling_expected, "sibling did not complete after removal" + assert request_b.status == og.RequestStatus.CLOSED def test_engine_teardown_and_recreation(model): @@ -403,15 +412,17 @@ def test_engine_teardown_and_recreation(model): first = og.Engine(model) sink1 = _Sink() - _add_request(first, model, _PROMPT_A, max_new, sink1) + first_request = _add_request(first, model, _PROMPT_A, max_new, sink1) _run(first) assert sink1.tokens == expected + assert first_request.status == og.RequestStatus.CLOSED del first gc.collect() second = og.Engine(model) assert not second.has_pending_requests() sink2 = _Sink() - _add_request(second, model, _PROMPT_A, max_new, sink2) + second_request = _add_request(second, model, _PROMPT_A, max_new, sink2) _run(second) assert sink2.tokens == expected + assert second_request.status == og.RequestStatus.CLOSED From 3aab3de1d8fae93e0e3bf8201e3b54c6bed5cbea Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 20:39:21 -0500 Subject: [PATCH 09/19] Harden Engine admission atomicity Prepare request, scheduler, and tracking state before publishing Engine ownership so failed admission stays retryable without duplicating the prompt. Files changed: - src/engine/engine.cpp - src/engine/request.cpp - src/engine/request.h - src/engine/scheduler.cpp - src/engine/scheduler.h - test/engine/request_lifecycle_tests.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/engine/engine.cpp | 20 ++--- src/engine/request.cpp | 49 +++++++++-- src/engine/request.h | 28 ++++-- src/engine/scheduler.cpp | 45 ++++++++-- src/engine/scheduler.h | 23 ++++- test/engine/request_lifecycle_tests.cpp | 112 ++++++++++++++++++++++++ 6 files changed, 241 insertions(+), 36 deletions(-) diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp index 1e56b66ecc..f9ad74f2e4 100644 --- a/src/engine/engine.cpp +++ b/src/engine/engine.cpp @@ -60,17 +60,15 @@ void Engine::AddRequest(std::shared_ptr 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(); + auto scheduler_preparation = scheduler_->PrepareAddRequest(request); + tracked_requests_.reserve(tracked_requests_.size() + 1); + + 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) { diff --git a/src/engine/request.cpp b/src/engine/request.cpp index c352b3818a..d2c7e88e4c 100644 --- a/src/engine/request.cpp +++ b/src/engine/request.cpp @@ -9,6 +9,13 @@ namespace Generators { +RequestAdmissionPreparation::RequestAdmissionPreparation() = default; +RequestAdmissionPreparation::~RequestAdmissionPreparation() = default; +RequestAdmissionPreparation::RequestAdmissionPreparation( + RequestAdmissionPreparation&&) noexcept = default; +RequestAdmissionPreparation& RequestAdmissionPreparation::operator=( + RequestAdmissionPreparation&&) noexcept = default; + namespace { DeviceSpan AllocateOnDevice(GeneratorParams& params, @@ -66,24 +73,44 @@ bool Request::IsExternallyAbandoned() const noexcept { } void Request::Assign(std::shared_ptr engine) { + auto preparation = PrepareAdmission(); + CommitAdmission(std::move(engine), std::move(preparation)); +} + +RequestAdmissionPreparation Request::PrepareAdmission() const { if (status_ != RequestStatus::Unassigned) { throw std::runtime_error("Cannot add the request to the engine since it is already assigned."); } if (prefill_input_ids_.empty()) { throw std::runtime_error("Cannot add a request with no input tokens to the engine."); } - engine_ = engine; - status_ = RequestStatus::Assigned; + RequestAdmissionPreparation preparation; + preparation.search = CreateSearch(*params_); + preparation.search->DeferCompletion(true); auto device_tokens = AllocateOnDevice(*params_, prefill_input_ids_); + preparation.search->AppendTokens(device_tokens); + preparation.prompt_sequence_length = preparation.search->GetSequenceLength(); + preparation.seen_sequence_length = preparation.prompt_sequence_length; + preparation.tokens_host.reserve(params_->search.max_length); + preparation.unseen_token_indices.reserve(params_->search.max_length); + preparation.tokens_host.insert( + preparation.tokens_host.end(), prefill_input_ids_.begin(), prefill_input_ids_.end()); + return preparation; +} + +void Request::CommitAdmission(std::shared_ptr engine, + RequestAdmissionPreparation&& preparation) noexcept { + search_ = std::move(preparation.search); + batched_sampler_state_ = std::move(preparation.sampling_state); + tokens_host_ = std::move(preparation.tokens_host); + unseen_token_indices_ = std::move(preparation.unseen_token_indices); + prompt_sequence_length_ = preparation.prompt_sequence_length; + seen_sequence_length_ = preparation.seen_sequence_length; processed_sequence_length_ = 0; - search_->AppendTokens(device_tokens); - prompt_sequence_length_ = CurrentSequenceLength(); - seen_sequence_length_ = CurrentSequenceLength(); - tokens_host_.reserve(params_->search.max_length); - unseen_token_indices_.reserve(params_->search.max_length); - tokens_host_.insert(tokens_host_.end(), prefill_input_ids_.begin(), prefill_input_ids_.end()); prefill_input_ids_.clear(); + engine_ = std::move(engine); + status_ = RequestStatus::Assigned; } void Request::Schedule() { @@ -435,6 +462,12 @@ BatchedSamplerState& Request::SamplingState(BatchedSampler& sampler) { return *batched_sampler_state_; } +void Request::CommitSamplingState(std::unique_ptr state) noexcept { + if (state) { + batched_sampler_state_ = std::move(state); + } +} + void Request::CompleteGeneration() { search_->CompleteGeneration(); diff --git a/src/engine/request.h b/src/engine/request.h index 6ebe15565d..a81023af99 100644 --- a/src/engine/request.h +++ b/src/engine/request.h @@ -29,6 +29,22 @@ struct RequestStepResult { bool done{}; }; +struct RequestAdmissionPreparation { + RequestAdmissionPreparation(); + ~RequestAdmissionPreparation(); + RequestAdmissionPreparation(RequestAdmissionPreparation&&) noexcept; + RequestAdmissionPreparation& operator=(RequestAdmissionPreparation&&) noexcept; + RequestAdmissionPreparation(const RequestAdmissionPreparation&) = delete; + RequestAdmissionPreparation& operator=(const RequestAdmissionPreparation&) = delete; + + std::unique_ptr search; + std::unique_ptr sampling_state; + std::vector tokens_host; + std::vector unseen_token_indices; + int64_t prompt_sequence_length{}; + int64_t seen_sequence_length{}; +}; + /** * @class Request * @brief Manages the state and lifecycle of a user request within the engine. @@ -48,13 +64,12 @@ struct Request : std::enable_shared_from_this, */ Request(std::shared_ptr params); - /** - * @brief Assigns this request to a specific engine for processing. - * @param engine Shared pointer to the Engine to be used for processing this request. - * - * Once assigned, the request will finalize the prefill tokens and prepare for scheduling. - */ + // Compatibility helper for tests and direct scheduler clients. Engine admission uses the + // prepare/commit methods below so no ownership is published until every throwing step succeeds. void Assign(std::shared_ptr engine); + RequestAdmissionPreparation PrepareAdmission() const; + void CommitAdmission(std::shared_ptr engine, + RequestAdmissionPreparation&& preparation) noexcept; /** * @brief Updates the status of the request to Active and prepares it for processing. @@ -269,6 +284,7 @@ struct Request : std::enable_shared_from_this, * @brief Returns this request's persistent random state for the given batched sampler. */ BatchedSamplerState& SamplingState(BatchedSampler& sampler); + void CommitSamplingState(std::unique_ptr state) noexcept; /** * @brief Retrieves the generator parameters associated with this request. diff --git a/src/engine/scheduler.cpp b/src/engine/scheduler.cpp index 1ed3c99f8b..31440a052f 100644 --- a/src/engine/scheduler.cpp +++ b/src/engine/scheduler.cpp @@ -28,19 +28,36 @@ ScheduledRequests Scheduler::CreateScheduledRequests(const StepPlan& plan) { GetBatchedSamplingPlan()}; } +void Scheduler::AddRequest(std::shared_ptr request) { + auto preparation = PrepareAddRequest(request); + CommitAddRequest(std::move(request), std::move(preparation)); +} + StaticBatchScheduler::StaticBatchScheduler(std::shared_ptr model, std::shared_ptr cache_manager) : Scheduler{model}, model_{model}, cache_manager_{cache_manager} {} -void StaticBatchScheduler::AddRequest(std::shared_ptr request) { +SchedulerAdmissionPreparation StaticBatchScheduler::PrepareAddRequest( + const std::shared_ptr& request) { // The static batch decoder rebuilds its contiguous cache from the whole sequence every step, so it // cannot resume a half written prompt. Only the paged cache can hold one. if (request->SearchOptions().chunk_size.value_or(0) != 0) { throw std::runtime_error( "search.chunk_size requires dynamic batching; the static batch scheduler cannot chunk a prefill."); } - if (auto* sampler = GetBatchedSampler()) - request->SamplingState(*sampler); - requests_pool_.push_back(request); + SchedulerAdmissionPreparation preparation; + if (auto* sampler = GetBatchedSampler()) { + preparation.sampling_state = + sampler->CreateState(request->SearchOptions().random_seed); + } + requests_pool_.reserve(requests_pool_.size() + 1); + return preparation; +} + +void StaticBatchScheduler::CommitAddRequest( + std::shared_ptr request, + SchedulerAdmissionPreparation&& preparation) noexcept { + request->CommitSamplingState(std::move(preparation.sampling_state)); + requests_pool_.push_back(std::move(request)); } void StaticBatchScheduler::RemoveRequest(std::shared_ptr request) { @@ -113,10 +130,22 @@ bool StaticBatchScheduler::HasPendingRequests() const { DynamicBatchScheduler::DynamicBatchScheduler(std::shared_ptr model, std::shared_ptr cache_manager) : Scheduler{model}, model_{model}, cache_manager_{cache_manager} {} -void DynamicBatchScheduler::AddRequest(std::shared_ptr request) { - if (auto* sampler = GetBatchedSampler()) - request->SamplingState(*sampler); - requests_pool_.push_back(request); +SchedulerAdmissionPreparation DynamicBatchScheduler::PrepareAddRequest( + const std::shared_ptr& request) { + SchedulerAdmissionPreparation preparation; + if (auto* sampler = GetBatchedSampler()) { + preparation.sampling_state = + sampler->CreateState(request->SearchOptions().random_seed); + } + requests_pool_.reserve(requests_pool_.size() + 1); + return preparation; +} + +void DynamicBatchScheduler::CommitAddRequest( + std::shared_ptr request, + SchedulerAdmissionPreparation&& preparation) noexcept { + request->CommitSamplingState(std::move(preparation.sampling_state)); + requests_pool_.push_back(std::move(request)); } void DynamicBatchScheduler::RemoveRequest(std::shared_ptr request) { diff --git a/src/engine/scheduler.h b/src/engine/scheduler.h index 6d76fc9060..566b45eb1c 100644 --- a/src/engine/scheduler.h +++ b/src/engine/scheduler.h @@ -16,6 +16,10 @@ namespace Generators { +struct SchedulerAdmissionPreparation { + std::unique_ptr sampling_state; +}; + struct Scheduler { /** * @brief Constructs a Scheduler instance with the specified model and cache manager. @@ -33,7 +37,12 @@ struct Scheduler { * This function adds the request to the internal pool of requests and marks it * as pending for scheduling. */ - virtual void AddRequest(std::shared_ptr request) = 0; + void AddRequest(std::shared_ptr request); + virtual SchedulerAdmissionPreparation PrepareAddRequest( + const std::shared_ptr& request) = 0; + virtual void CommitAddRequest( + std::shared_ptr request, + SchedulerAdmissionPreparation&& preparation) noexcept = 0; /** * @brief Removes a request from the Scheduler. @@ -82,7 +91,11 @@ struct Scheduler { struct StaticBatchScheduler : Scheduler { StaticBatchScheduler(std::shared_ptr model, std::shared_ptr cache_manager); - void AddRequest(std::shared_ptr request) override; + SchedulerAdmissionPreparation PrepareAddRequest( + const std::shared_ptr& request) override; + void CommitAddRequest( + std::shared_ptr request, + SchedulerAdmissionPreparation&& preparation) noexcept override; void RemoveRequest(std::shared_ptr request) override; @@ -99,7 +112,11 @@ struct StaticBatchScheduler : Scheduler { struct DynamicBatchScheduler : Scheduler { DynamicBatchScheduler(std::shared_ptr model, std::shared_ptr cache_manager); - void AddRequest(std::shared_ptr request) override; + SchedulerAdmissionPreparation PrepareAddRequest( + const std::shared_ptr& request) override; + void CommitAddRequest( + std::shared_ptr request, + SchedulerAdmissionPreparation&& preparation) noexcept override; void RemoveRequest(std::shared_ptr request) override; diff --git a/test/engine/request_lifecycle_tests.cpp b/test/engine/request_lifecycle_tests.cpp index e0d8dd33e0..6cb2a10a7b 100644 --- a/test/engine/request_lifecycle_tests.cpp +++ b/test/engine/request_lifecycle_tests.cpp @@ -33,6 +33,62 @@ DeviceSpan LogitsForToken(Model& model, int32_t token) { return logits; } +class FailingAllocationDevice final : public DeviceInterface { + public: + explicit FailingAllocationDevice(DeviceInterface& inner) : inner_{inner} {} + + DeviceType GetType() const override { return inner_.GetType(); } + void InitOrt(const OrtApi& api, Ort::Allocator& allocator) override { + inner_.InitOrt(api, allocator); + } + Ort::Allocator& GetAllocator() override { return inner_.GetAllocator(); } + std::unique_ptr GetMemoryInfo() const override { + return inner_.GetMemoryInfo(); + } + std::shared_ptr AllocateBase(size_t size) override { + if (fail_allocation_) { + throw std::runtime_error("Injected request preparation allocation failure."); + } + return inner_.AllocateBase(size); + } + std::shared_ptr WrapMemoryBase(void* memory, size_t size) override { + return inner_.WrapMemoryBase(memory, size); + } + std::unique_ptr CreateGreedy(const GeneratorParams& params) override { + return inner_.CreateGreedy(params); + } + std::unique_ptr CreateBeam(const GeneratorParams& params) override { + return inner_.CreateBeam(params); + } + void Synchronize() override { inner_.Synchronize(); } + + void SetFailAllocation(bool fail) { fail_allocation_ = fail; } + + private: + DeviceInterface& inner_; + bool fail_allocation_{true}; +}; + +class FailingAdmissionScheduler final : public DynamicBatchScheduler { + public: + FailingAdmissionScheduler(std::shared_ptr model, + std::shared_ptr cache_manager) + : DynamicBatchScheduler(std::move(model), std::move(cache_manager)) {} + + SchedulerAdmissionPreparation PrepareAddRequest( + const std::shared_ptr& request) override { + if (fail_preparation_) { + throw std::runtime_error("Injected scheduler admission preparation failure."); + } + return DynamicBatchScheduler::PrepareAddRequest(request); + } + + void SetFailPreparation(bool fail) { fail_preparation_ = fail; } + + private: + bool fail_preparation_{true}; +}; + class RequestLifecycleTest : public ::testing::Test { protected: void SetUp() override { @@ -69,6 +125,62 @@ TEST_F(RequestLifecycleTest, EmptyRequestIsRejectedBeforeAssignment) { EXPECT_EQ(request->status_, RequestStatus::Unassigned); } +TEST_F(RequestLifecycleTest, + DevicePreparationFailureLeavesAdmissionUnassignedAndRetryAppendsPromptOnce) { + const auto prompt = Prompt(); + auto request = NewRequest(); + request->AddTokens(prompt); + auto params = request->Params(); + FailingAllocationDevice failing_device{*params->p_device}; + params->p_device = &failing_device; + + EXPECT_THROW(engine_.engine->AddRequest(request), std::runtime_error); + EXPECT_EQ(request->Status(), RequestStatus::Unassigned); + EXPECT_FALSE(engine_.engine->HasPendingRequests()); + EXPECT_EQ(engine_.cache->AllocatedCount(), 0u); + + failing_device.SetFailAllocation(false); + EXPECT_NO_THROW(engine_.engine->AddRequest(request)); + EXPECT_EQ(request->Status(), RequestStatus::Assigned); + EXPECT_EQ(request->CurrentSequenceLength(), + static_cast(prompt.size())); + EXPECT_TRUE(engine_.engine->HasPendingRequests()); + EXPECT_EQ(engine_.cache->AllocatedCount(), 0u); + params->p_device = model_->p_device_scoring_; +} + +TEST_F(RequestLifecycleTest, + SchedulerPreparationFailureLeavesAdmissionUnassignedAndRetryAppendsPromptOnce) { + const auto prompt = Prompt(); + auto request = NewRequest(); + request->AddTokens(prompt); + + auto cache = + std::make_shared(model_, /*capacity=*/8); + auto scheduler = + std::make_unique(model_, cache); + auto* scheduler_observer = scheduler.get(); + auto executor = std::make_unique( + model_, cache, EosToken(*model_)); + EngineDependencies dependencies{ + cache, std::move(scheduler), std::move(executor)}; + auto engine = + std::make_shared(model_, std::move(dependencies)); + + EXPECT_THROW(engine->AddRequest(request), std::runtime_error); + EXPECT_EQ(request->Status(), RequestStatus::Unassigned); + EXPECT_FALSE(engine->HasPendingRequests()); + EXPECT_EQ(cache->AllocatedCount(), 0u); + + scheduler_observer->SetFailPreparation(false); + EXPECT_NO_THROW(engine->AddRequest(request)); + EXPECT_EQ(request->Status(), RequestStatus::Assigned); + EXPECT_EQ(request->CurrentSequenceLength(), + static_cast(prompt.size())); + EXPECT_TRUE(engine->HasPendingRequests()); + EXPECT_EQ(cache->AllocatedCount(), 0u); +} + // An append that would exceed the model's max length is rejected before any tokens are buffered, so // a subsequent valid append and assign reflect only the accepted tokens. TEST_F(RequestLifecycleTest, AppendBeyondContextIsRejectedBeforeMutation) { From 95734542f791ca7388bb91501ce76703dfe957d0 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 20:39:31 -0500 Subject: [PATCH 10/19] Strengthen paged cache diagnostics Cross-check request progress with full and window cache ownership, including attributed reservation blocks, so diagnostic snapshots expose inconsistent committed boundaries. Files changed: - src/engine/engine_invariants.cpp - src/engine/engine_invariants.h - src/engine/paged_key_value_cache.cpp - test/engine/engine_invariants_tests.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/engine/engine_invariants.cpp | 88 ++++++++++++++ src/engine/engine_invariants.h | 2 + src/engine/paged_key_value_cache.cpp | 9 ++ test/engine/engine_invariants_tests.cpp | 146 +++++++++++++++++++++++- 4 files changed, 241 insertions(+), 4 deletions(-) diff --git a/src/engine/engine_invariants.cpp b/src/engine/engine_invariants.cpp index 10552265a9..adeaf121f4 100644 --- a/src/engine/engine_invariants.cpp +++ b/src/engine/engine_invariants.cpp @@ -115,6 +115,7 @@ std::vector ValidateCacheInvariants(const PagedCacheSnapshot std::unordered_set reservation_requests; std::unordered_set blocks_assigned_to_delta; + std::unordered_set window_blocks_assigned_to_delta; for (const auto& reservation : cache.reservations) { if (!reservation_requests.insert(reservation.request_id).second) { add("Request " + PtrId(reservation.request_id) + @@ -129,6 +130,20 @@ std::vector ValidateCacheInvariants(const PagedCacheSnapshot add("Request " + PtrId(reservation.request_id) + " transaction tail-slot growth exceeds total growth."); } + const auto committed_owner = std::find_if( + cache.requests.begin(), cache.requests.end(), + [&reservation](const RequestBlockSnapshot& request) { + return request.request_id == reservation.request_id; + }); + if (reservation.newly_admitted == (committed_owner != cache.requests.end())) { + add("Request " + PtrId(reservation.request_id) + + " transaction membership disagrees with committed cache ownership."); + } + if (committed_owner != cache.requests.end() && + reservation.committed_slots != committed_owner->used_slots) { + add("Request " + PtrId(reservation.request_id) + + " transaction committed slots disagree with committed cache usage."); + } for (const size_t block_id : reservation.reserved_block_ids) { if (reserved_blocks.find(block_id) == reserved_blocks.end()) { add("Request " + PtrId(reservation.request_id) + @@ -140,6 +155,12 @@ std::vector ValidateCacheInvariants(const PagedCacheSnapshot " is assigned to more than one Request delta."); } } + for (const size_t block_id : reservation.reserved_window_block_ids) { + if (!window_blocks_assigned_to_delta.insert(block_id).second) { + add("Transaction-reserved window block id " + std::to_string(block_id) + + " is assigned to more than one Request delta."); + } + } } if (blocks_assigned_to_delta != reserved_blocks) { add("Not every transaction-reserved block belongs to exactly one Request delta."); @@ -187,6 +208,25 @@ std::vector ValidateCacheInvariants(const PagedCacheSnapshot " is also committed to a Request."); } } + if (window_blocks_assigned_to_delta != reserved_window_blocks) { + add("Not every transaction-reserved window block belongs to exactly one Request delta."); + } + if (window.total_blocks == 0 && !window_blocks_assigned_to_delta.empty()) { + add("A non-windowed cache has Request-attributed window reservations."); + } + if (window.total_blocks != 0) { + for (const auto& reservation : cache.reservations) { + const size_t expected_window_blocks = + reservation.newly_admitted ? window.blocks_per_request : 0; + if (reservation.reserved_window_block_ids.size() != + expected_window_blocks) { + add("Request " + PtrId(reservation.request_id) + " owns " + + std::to_string(reservation.reserved_window_block_ids.size()) + + " transaction-reserved window blocks instead of " + + std::to_string(expected_window_blocks) + "."); + } + } + } if (window.free_blocks > window.total_blocks) { add("window free_blocks (" + std::to_string(window.free_blocks) + ") exceeds total_blocks (" + std::to_string(window.total_blocks) + ")."); @@ -255,6 +295,54 @@ std::vector ValidateInvariants(const PagedCacheSnapshot& cac "Cache holds a block table for unknown Request " + PtrId(owner.request_id) + "."}); } } + for (const auto& owner : cache.window_blocks.requests) { + if (known_requests.find(owner.request_id) == known_requests.end()) { + violations.push_back(InvariantViolation{ + "Window cache holds a block table for unknown Request " + + PtrId(owner.request_id) + "."}); + } + } + for (const auto& reservation : cache.reservations) { + if (known_requests.find(reservation.request_id) == known_requests.end()) { + violations.push_back(InvariantViolation{ + "Cache holds a transaction reservation for unknown Request " + + PtrId(reservation.request_id) + "."}); + } + } + + std::unordered_map processed_by_request; + for (const auto& request : requests) { + if (request.processed_sequence_length >= 0) { + processed_by_request.emplace( + request.request_id, + static_cast(request.processed_sequence_length)); + } + } + for (const auto& owner : cache.requests) { + const auto processed = processed_by_request.find(owner.request_id); + if (processed != processed_by_request.end() && + owner.used_slots != processed->second) { + violations.push_back(InvariantViolation{ + "Request " + PtrId(owner.request_id) + " committed cache usage (" + + std::to_string(owner.used_slots) + ") differs from processed sequence length (" + + std::to_string(processed->second) + ")."}); + } + } + + if (cache.window_blocks.total_blocks != 0) { + std::set full_owners; + std::set window_owners; + for (const auto& owner : cache.requests) { + full_owners.insert(owner.request_id); + } + for (const auto& owner : cache.window_blocks.requests) { + window_owners.insert(owner.request_id); + } + if (full_owners != window_owners) { + violations.push_back(InvariantViolation{ + "Full-cache and window-cache owner sets disagree."}); + } + } return violations; } diff --git a/src/engine/engine_invariants.h b/src/engine/engine_invariants.h index 164c56ed81..d69ed97294 100644 --- a/src/engine/engine_invariants.h +++ b/src/engine/engine_invariants.h @@ -53,6 +53,8 @@ struct RequestReservationSnapshot { size_t target_slots{}; size_t tail_slots_to_consume{}; std::vector reserved_block_ids; + std::vector reserved_window_block_ids; + bool newly_admitted{}; }; struct WindowBlockPoolSnapshot { diff --git a/src/engine/paged_key_value_cache.cpp b/src/engine/paged_key_value_cache.cpp index f603393a46..bd29a98fd9 100644 --- a/src/engine/paged_key_value_cache.cpp +++ b/src/engine/paged_key_value_cache.cpp @@ -499,6 +499,15 @@ PagedCacheSnapshot PagedKeyValueCache::Snapshot( request_reservation.reserved_block_ids.push_back( reservation.ReservedBlocks()[delta.reserved_block_offset + i]->Id()); } + request_reservation.reserved_window_block_ids.reserve( + delta.reserved_window_block_count); + for (size_t i = 0; i < delta.reserved_window_block_count; ++i) { + request_reservation.reserved_window_block_ids.push_back( + reservation.ReservedWindowBlocks()[ + delta.reserved_window_block_offset + i] + ->Id()); + } + request_reservation.newly_admitted = delta.newly_admitted; snapshot.reservations.push_back(std::move(request_reservation)); } return snapshot; diff --git a/test/engine/engine_invariants_tests.cpp b/test/engine/engine_invariants_tests.cpp index 5e9ff01650..4c799b6e46 100644 --- a/test/engine/engine_invariants_tests.cpp +++ b/test/engine/engine_invariants_tests.cpp @@ -102,7 +102,14 @@ TEST(InvariantValidatorTest, WindowBlockPoolValidatesRingSizeAndOwnershipIndepen } TEST(InvariantValidatorTest, WindowBlockPoolIncludesTransactionReservationsInAccounting) { - auto cache = MakeValidCache(); + PagedCacheSnapshot cache; + cache.block_size = kBlockSize; + cache.total_blocks = 1; + cache.transaction_reserved_block_ids = {0}; + cache.reservations = { + RequestReservationSnapshot{ + kRequestA, 0, 1, 0, {0}, {0, 1}, true}, + }; cache.window_blocks.total_blocks = 2; cache.window_blocks.free_blocks = 0; cache.window_blocks.blocks_per_request = 2; @@ -148,7 +155,7 @@ TEST(InvariantValidatorTest, InitialAdmissionReservationValidatesWithoutCommitte cache.free_blocks = 0; cache.transaction_reserved_block_ids = {0}; cache.reservations = { - RequestReservationSnapshot{kRequestA, 0, 1, 0, {0}}, + RequestReservationSnapshot{kRequestA, 0, 1, 0, {0}, {}, true}, }; EXPECT_TRUE(ValidateCacheInvariants(cache).empty()); @@ -176,7 +183,7 @@ TEST(InvariantValidatorTest, InitialAdmissionRejectsUnreservedDeltaBlock) { cache.free_blocks = 1; cache.transaction_reserved_block_ids = {0}; cache.reservations = { - RequestReservationSnapshot{kRequestA, 0, 1, 0, {1}}, + RequestReservationSnapshot{kRequestA, 0, 1, 0, {1}, {}, true}, }; const auto violations = ValidateCacheInvariants(cache); @@ -329,7 +336,7 @@ TEST(InvariantValidatorTest, TurnCompleteRequestFullyProcessedIsValid) { TEST(InvariantValidatorTest, ConsistentSnapshotsValidateClean) { const auto cache = MakeValidCache(); const std::vector requests{ - MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 5, 9), MakeValidRequest(kRequestB, RequestStatus::Active, 4, 4, 4), }; EXPECT_TRUE(ValidateInvariants(cache, requests).empty()); @@ -345,6 +352,137 @@ TEST(InvariantValidatorTest, BlockTableForUnknownRequestReported) { EXPECT_FALSE(ValidateInvariants(cache, requests).empty()); } +TEST(InvariantValidatorTest, CommittedCacheUsageMustMatchProcessedLength) { + const auto cache = MakeValidCache(); + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 8, 8), + MakeValidRequest(kRequestB, RequestStatus::Active, 4, 4, 4), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "differs from processed sequence length") != + std::string::npos; + }), + violations.end()); +} + +TEST(InvariantValidatorTest, FullAndWindowOwnerSetsMustAgree) { + auto cache = MakeValidCache(); + cache.window_blocks.total_blocks = 2; + cache.window_blocks.blocks_per_request = 2; + cache.window_blocks.requests = { + RequestBlockSnapshot{kRequestA, {0, 1}}, + }; + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + MakeValidRequest(kRequestB, RequestStatus::Active, 4, 4, 4), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "Full-cache and window-cache owner sets disagree") != + std::string::npos; + }), + violations.end()); +} + +TEST(InvariantValidatorTest, WindowOwnerMustBeAKnownRequest) { + auto cache = MakeValidCache(); + cache.window_blocks.total_blocks = 4; + cache.window_blocks.blocks_per_request = 2; + cache.window_blocks.requests = { + RequestBlockSnapshot{kRequestA, {0, 1}}, + RequestBlockSnapshot{kRequestB, {2, 3}}, + }; + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "Window cache holds a block table for unknown Request") != + std::string::npos; + }), + violations.end()); +} + +TEST(InvariantValidatorTest, ReservationOwnerMustBeKnownAndMatchCommittedUsage) { + auto cache = MakeValidCache(); + cache.free_blocks = 0; + cache.transaction_reserved_block_ids = {3}; + cache.reservations = { + RequestReservationSnapshot{ + kRequestB, /*committed_slots=*/3, /*target_slots=*/5, + /*tail_slots_to_consume=*/0, /*reserved_block_ids=*/{3}}, + }; + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "transaction committed slots disagree") != + std::string::npos; + }), + violations.end()); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "transaction reservation for unknown Request") != + std::string::npos; + }), + violations.end()); +} + +TEST(InvariantValidatorTest, WindowReservationMustHaveConsistentOwnership) { + PagedCacheSnapshot cache; + cache.block_size = kBlockSize; + cache.total_blocks = 1; + cache.transaction_reserved_block_ids = {0}; + cache.window_blocks.total_blocks = 2; + cache.window_blocks.blocks_per_request = 2; + cache.window_blocks.transaction_reserved_block_ids = {0, 1}; + cache.reservations = { + RequestReservationSnapshot{ + kRequestA, 0, 1, 0, {0}, {0}, true}, + }; + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Assigned, 3, 0, 3), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "Not every transaction-reserved window block belongs") != + std::string::npos; + }), + violations.end()); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "transaction-reserved window blocks instead of 2") != + std::string::npos; + }), + violations.end()); +} + TEST(InvariantValidatorTest, ThrowWrapperListsViolations) { auto cache = MakeValidCache(); cache.free_blocks = 0; // break block accounting From 41a83078cd7b258fb0ed03ea291661c32ac77954 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 20:39:41 -0500 Subject: [PATCH 11/19] Cover staged Search rollback after ring writes Inject a retryable fault after real windowed cache writes and Search staging, then prove rollback and clean-replay parity while documenting the hardened admission and diagnostic contracts. Files changed: - test/engine/windowed_transaction_tests.cpp - docs/paged_attention_engine.md Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/paged_attention_engine.md | 20 ++- test/engine/windowed_transaction_tests.cpp | 147 +++++++++++++++++++-- 2 files changed, 154 insertions(+), 13 deletions(-) diff --git a/docs/paged_attention_engine.md b/docs/paged_attention_engine.md index 2da0682f64..6f1a12c9eb 100644 --- a/docs/paged_attention_engine.md +++ b/docs/paged_attention_engine.md @@ -61,6 +61,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: @@ -113,15 +120,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. diff --git a/test/engine/windowed_transaction_tests.cpp b/test/engine/windowed_transaction_tests.cpp index 4c6805b08f..3b88a34558 100644 --- a/test/engine/windowed_transaction_tests.cpp +++ b/test/engine/windowed_transaction_tests.cpp @@ -152,22 +152,35 @@ class ObservingPagedCacheManager final : public PagedCacheManager { // and raise a retryable failure. class RetryableAfterRealDecodeExecutor final : public ModelExecutor { public: + enum class FailurePoint { + None, + BeforeSearchMutation, + AfterSearchMutation, + }; + explicit RetryableAfterRealDecodeExecutor( - std::unique_ptr inner) - : inner_{std::move(inner)} {} + std::unique_ptr inner, + std::shared_ptr cache) + : inner_{std::move(inner)}, cache_{std::move(cache)} {} void Decode(ScheduledRequests& scheduled_requests, ExecutionContext& context) override { inner_->Decode(scheduled_requests, context); ++completed_real_decodes_; - if (!fail_next_decode_) { + auto failure_point = failure_point_; + if (failure_point == FailurePoint::AfterSearchMutation && + context.plan && !context.plan->requests.empty() && + context.plan->requests.front().target_cache_slots != + failure_target_cache_slots_) { + failure_point = FailurePoint::None; + } else { + failure_point_ = FailurePoint::None; + } + if (failure_point == FailurePoint::None) { return; } - fail_next_decode_ = false; - const auto logits = scheduled_requests.ProcessLogits(); - observed_real_logits_ = logits.size() == scheduled_requests.size(); if (!context.plan || context.plan->requests.size() != 1) { throw std::logic_error( "Windowed rollback fault expected one planned request."); @@ -175,6 +188,23 @@ class RetryableAfterRealDecodeExecutor final : public ModelExecutor { const auto& entry = context.plan->requests.front(); failed_unprocessed_token_count_ = entry.unprocessed_token_count; failed_target_cache_slots_ = entry.target_cache_slots; + if (failure_point == FailurePoint::AfterSearchMutation) { + request_before_failure_ = entry.request->Snapshot(); + cache_before_failure_ = cache_->Snapshot(); + std::vector staged_results; + scheduled_requests.GenerateNextTokensForTransaction( + *context.plan, staged_results); + observed_real_logits_ = + staged_results.size() == scheduled_requests.size(); + staged_sequence_length_ = + entry.request->CurrentSequenceLength(); + staged_search_mutation_ = + staged_sequence_length_ > entry.sequence_length_before; + } else { + const auto logits = scheduled_requests.ProcessLogits(); + observed_real_logits_ = + logits.size() == scheduled_requests.size(); + } injected_failure_ = true; throw ModelExecutionError{ ExecutionFailureKind::RetryableAbort, @@ -182,7 +212,14 @@ class RetryableAfterRealDecodeExecutor final : public ModelExecutor { }; } - void FailNextDecode() { fail_next_decode_ = true; } + void FailNextDecode() { + failure_point_ = FailurePoint::BeforeSearchMutation; + } + void FailNextDecodeAfterSearchMutationAtTarget( + size_t target_cache_slots) { + failure_point_ = FailurePoint::AfterSearchMutation; + failure_target_cache_slots_ = target_cache_slots; + } bool InjectedFailure() const { return injected_failure_; } bool ObservedRealLogits() const { return observed_real_logits_; } @@ -193,15 +230,31 @@ class RetryableAfterRealDecodeExecutor final : public ModelExecutor { size_t FailedTargetCacheSlots() const { return failed_target_cache_slots_; } + bool StagedSearchMutation() const { return staged_search_mutation_; } + int64_t StagedSequenceLength() const { + return staged_sequence_length_; + } + const RequestStateSnapshot& RequestBeforeFailure() const { + return request_before_failure_; + } + const PagedCacheSnapshot& CacheBeforeFailure() const { + return cache_before_failure_; + } private: std::unique_ptr inner_; - bool fail_next_decode_{}; + std::shared_ptr cache_; + FailurePoint failure_point_{FailurePoint::None}; + size_t failure_target_cache_slots_{}; bool injected_failure_{}; bool observed_real_logits_{}; + bool staged_search_mutation_{}; size_t completed_real_decodes_{}; size_t failed_unprocessed_token_count_{}; size_t failed_target_cache_slots_{}; + int64_t staged_sequence_length_{}; + RequestStateSnapshot request_before_failure_; + PagedCacheSnapshot cache_before_failure_; }; struct FaultInjectingEngine { @@ -216,7 +269,7 @@ FaultInjectingEngine MakeFaultInjectingEngine( std::make_shared(model); auto scheduler = Scheduler::Create(model, cache); auto executor = std::make_unique( - ModelExecutor::Create(model, cache)); + ModelExecutor::Create(model, cache), cache); auto* executor_observer = executor.get(); EngineDependencies dependencies{ cache, std::move(scheduler), std::move(executor)}; @@ -380,6 +433,82 @@ TEST(WindowedTransactionTest, faulting.engine->RemoveRequest(request); } +TEST(WindowedTransactionTest, + ContinuedRingWritesAndStagedSearchRollbackTogetherBeforeRetry) { + const auto clean_replay_output = RunCleanReplay(); + ASSERT_EQ(clean_replay_output, kSecondTurnOutput); + + auto model = LoadWindowedMultiwrapModel(); + auto faulting = MakeFaultInjectingEngine(model); + auto request = MintRequest(*model, kInitialPrompt); + faulting.engine->AddRequest(request); + ASSERT_EQ(RunTurn(faulting.engine, request), kFirstTurnOutput); + request->Continue(kContinuation); + + ASSERT_EQ(request->Status(), RequestStatus::Assigned); + ASSERT_FALSE(request->HasUnseenTokens()); + + const size_t reservations_before = + faulting.cache->ReservationCount(); + const size_t commits_before = faulting.cache->CommitCount(); + const size_t releases_before = faulting.cache->ReleaseCount(); + // Step commits four two-token continuation chunks before its final one-token prefill. Inject only + // at target 12: that transaction writes absolute position 11 into ring slot 3 after multiple + // wraps, then stages token 28 in Search before raising the retryable failure. + faulting.executor->FailNextDecodeAfterSearchMutationAtTarget(12); + + try { + static_cast(faulting.engine->Step()); + FAIL() << "Expected retryable failure after staged Search mutation."; + } catch (const EngineStepError& error) { + EXPECT_EQ(error.Outcome().kind, + StepOutcomeKind::RetryableBatchAbort); + } + + EXPECT_TRUE(faulting.executor->InjectedFailure()); + EXPECT_TRUE(faulting.executor->ObservedRealLogits()); + EXPECT_TRUE(faulting.executor->StagedSearchMutation()); + EXPECT_EQ(faulting.executor->StagedSequenceLength(), 13); + EXPECT_EQ(faulting.executor->FailedUnprocessedTokenCount(), 1u); + EXPECT_EQ(faulting.executor->FailedTargetCacheSlots(), 12u); + + const auto& request_before = + faulting.executor->RequestBeforeFailure(); + const auto& cache_before = + faulting.executor->CacheBeforeFailure(); + ASSERT_EQ(request_before.status, RequestStatus::Active); + ASSERT_EQ(request_before.current_sequence_length, 12); + ASSERT_EQ(request_before.processed_sequence_length, 11); + + const auto request_after = request->Snapshot(); + EXPECT_EQ(request_after.status, request_before.status); + EXPECT_EQ(request_after.current_sequence_length, + request_before.current_sequence_length); + EXPECT_EQ(request_after.processed_sequence_length, + request_before.processed_sequence_length); + EXPECT_FALSE(request->HasUnseenTokens()); + EXPECT_TRUE(faulting.engine->HasPendingRequests()); + + const auto cache_after = faulting.cache->Snapshot(); + ExpectCacheOwnershipRestored(cache_after, cache_before); + EXPECT_EQ(faulting.cache->ReservationCount(), + reservations_before + 5); + EXPECT_EQ(faulting.cache->CommitCount(), commits_before + 4); + EXPECT_EQ(faulting.cache->ReleaseCount(), + releases_before + 1); + EXPECT_EQ(faulting.cache->ActiveReservations(), 0u); + EXPECT_NO_THROW(ThrowIfInvariantsViolated( + cache_after, std::vector{request_after})); + + const auto retry_output = RunTurn(faulting.engine, request); + EXPECT_EQ(retry_output, clean_replay_output); + EXPECT_EQ(retry_output, kSecondTurnOutput); + EXPECT_EQ(request->Status(), RequestStatus::TurnComplete); + EXPECT_EQ(request->CurrentSequenceLength(), 14); + + faulting.engine->RemoveRequest(request); +} + } // namespace } // namespace test } // namespace Generators From 14bcb1493ef216d47adb6f00a553adf7caf3858f Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 18 Aug 2026 21:12:50 -0500 Subject: [PATCH 12/19] Make sampler admission failure atomic Reserve Engine and scheduler storage before CUDA sampler creation, and keep sampler pool counters and free-list ownership unchanged when state growth or initialization fails. Files changed: - src/cuda/interface.cpp - src/cuda/sampler_state_index_pool.h - src/engine/engine.cpp - src/engine/scheduler.cpp - test/engine/sampler_state_index_pool_tests.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/cuda/interface.cpp | 37 +++++----- src/cuda/sampler_state_index_pool.h | 54 +++++++++++++++ src/engine/engine.cpp | 2 +- src/engine/scheduler.cpp | 4 +- .../engine/sampler_state_index_pool_tests.cpp | 67 +++++++++++++++++++ 5 files changed, 139 insertions(+), 25 deletions(-) create mode 100644 src/cuda/sampler_state_index_pool.h create mode 100644 test/engine/sampler_state_index_pool_tests.cpp diff --git a/src/cuda/interface.cpp b/src/cuda/interface.cpp index 528ab21176..efb61ca714 100644 --- a/src/cuda/interface.cpp +++ b/src/cuda/interface.cpp @@ -8,6 +8,7 @@ #include "search_cuda.h" #include "kernels.h" #include "cuda_topk.h" +#include "sampler_state_index_pool.h" #include #include #include @@ -93,25 +94,18 @@ 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(std::random_device{}()) - : static_cast(random_seed); - cuda::LaunchInitCurandState(seed, states_.Span().data() + index, GetStream()); - return index; + return indices_.Acquire([this, random_seed](int index, int required_size) { + EnsureCapacity(required_size); + const unsigned long long seed = + random_seed == -1 + ? static_cast(std::random_device{}()) + : static_cast(random_seed); + cuda::LaunchInitCurandState( + seed, states_.Span().data() + index, GetStream()); + }); } - void Release(int index) { - free_indices_.push_back(index); - } + void Release(int index) noexcept { indices_.Release(index); } curandState* Data() { return states_.Span().data(); } @@ -122,9 +116,9 @@ struct CudaSamplerStatePool { const int new_capacity = std::max(required_capacity, std::max(4, capacity_ * 2)); auto new_states = AllocateCudaSpan(new_capacity); - if (size_ > 1) { + if (indices_.Size() > 0) { CUDA_CHECK(cudaMemcpyAsync(new_states.Span().data(), states_.Span().data(), - static_cast(size_ - 1) * sizeof(curandState), + static_cast(indices_.Size()) * sizeof(curandState), cudaMemcpyDeviceToDevice, GetStream())); CUDA_CHECK(cudaStreamSynchronize(GetStream())); } @@ -133,8 +127,7 @@ struct CudaSamplerStatePool { } DeviceSpan states_; - std::vector free_indices_; - int size_{}; + SamplerStateIndexPool indices_; int capacity_{}; }; @@ -142,7 +135,7 @@ struct CudaBatchedSamplerState final : BatchedSamplerState { CudaBatchedSamplerState(std::shared_ptr pool, int index) : pool_{std::move(pool)}, index_{index} {} - ~CudaBatchedSamplerState() override { pool_->Release(index_); } + ~CudaBatchedSamplerState() noexcept override { pool_->Release(index_); } std::shared_ptr pool_; int index_{}; diff --git a/src/cuda/sampler_state_index_pool.h b/src/cuda/sampler_state_index_pool.h new file mode 100644 index 0000000000..aa3b429255 --- /dev/null +++ b/src/cuda/sampler_state_index_pool.h @@ -0,0 +1,54 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include + +namespace Generators { + +class SamplerStateIndexPool { + public: + template + int Acquire(Prepare&& prepare) { + 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 noexcept, so reserve its future slot before any external preparation can + // publish an acquired index. A throwing preparation leaves every pool counter unchanged. + free_indices_.reserve(static_cast(required_size)); + std::forward(prepare)(index, required_size); + + if (reusing) { + free_indices_.pop_back(); + } else { + size_ = required_size; + } + return index; + } + + 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_; } + size_t FreeCount() const noexcept { return free_indices_.size(); } + size_t ActiveCount() const noexcept { + return static_cast(size_) - free_indices_.size(); + } + + private: + std::vector free_indices_; + int size_{}; +}; + +} // namespace Generators diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp index f9ad74f2e4..4275e78277 100644 --- a/src/engine/engine.cpp +++ b/src/engine/engine.cpp @@ -61,8 +61,8 @@ void Engine::AddRequest(std::shared_ptr request) { } auto request_preparation = request->PrepareAdmission(); - auto scheduler_preparation = scheduler_->PrepareAddRequest(request); tracked_requests_.reserve(tracked_requests_.size() + 1); + auto scheduler_preparation = scheduler_->PrepareAddRequest(request); request_preparation.sampling_state = std::move(scheduler_preparation.sampling_state); diff --git a/src/engine/scheduler.cpp b/src/engine/scheduler.cpp index 31440a052f..aefee6954c 100644 --- a/src/engine/scheduler.cpp +++ b/src/engine/scheduler.cpp @@ -44,12 +44,12 @@ SchedulerAdmissionPreparation StaticBatchScheduler::PrepareAddRequest( throw std::runtime_error( "search.chunk_size requires dynamic batching; the static batch scheduler cannot chunk a prefill."); } + requests_pool_.reserve(requests_pool_.size() + 1); SchedulerAdmissionPreparation preparation; if (auto* sampler = GetBatchedSampler()) { preparation.sampling_state = sampler->CreateState(request->SearchOptions().random_seed); } - requests_pool_.reserve(requests_pool_.size() + 1); return preparation; } @@ -132,12 +132,12 @@ DynamicBatchScheduler::DynamicBatchScheduler(std::shared_ptr model, std:: SchedulerAdmissionPreparation DynamicBatchScheduler::PrepareAddRequest( const std::shared_ptr& request) { + requests_pool_.reserve(requests_pool_.size() + 1); SchedulerAdmissionPreparation preparation; if (auto* sampler = GetBatchedSampler()) { preparation.sampling_state = sampler->CreateState(request->SearchOptions().random_seed); } - requests_pool_.reserve(requests_pool_.size() + 1); return preparation; } diff --git a/test/engine/sampler_state_index_pool_tests.cpp b/test/engine/sampler_state_index_pool_tests.cpp new file mode 100644 index 0000000000..6bacda9eb2 --- /dev/null +++ b/test/engine/sampler_state_index_pool_tests.cpp @@ -0,0 +1,67 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include + +#include "cuda/sampler_state_index_pool.h" + +namespace Generators { +namespace { + +TEST(SamplerStateIndexPoolTest, NewIndexPreparationFailureIsAtomic) { + SamplerStateIndexPool pool; + + EXPECT_THROW( + pool.Acquire([](int index, int required_size) { + EXPECT_EQ(index, 0); + EXPECT_EQ(required_size, 1); + throw std::runtime_error("Injected sampler state growth failure."); + }), + std::runtime_error); + EXPECT_EQ(pool.Size(), 0); + EXPECT_EQ(pool.ActiveCount(), 0u); + EXPECT_EQ(pool.FreeCount(), 0u); + + const int index = + pool.Acquire([](int, int) {}); + EXPECT_EQ(index, 0); + EXPECT_EQ(pool.Size(), 1); + EXPECT_EQ(pool.ActiveCount(), 1u); +} + +TEST(SamplerStateIndexPoolTest, ReusedIndexPreparationFailureKeepsItFree) { + SamplerStateIndexPool pool; + const int original = pool.Acquire([](int, int) {}); + pool.Release(original); + + EXPECT_THROW( + pool.Acquire([](int, int) { + throw std::runtime_error("Injected sampler state initialization failure."); + }), + std::runtime_error); + EXPECT_EQ(pool.Size(), 1); + EXPECT_EQ(pool.ActiveCount(), 0u); + EXPECT_EQ(pool.FreeCount(), 1u); + + const int retried = pool.Acquire([](int, int) {}); + EXPECT_EQ(retried, original); + EXPECT_EQ(pool.ActiveCount(), 1u); +} + +TEST(SamplerStateIndexPoolTest, ReleaseIsNoexceptAndAllocationFree) { + SamplerStateIndexPool pool; + const int first = pool.Acquire([](int, int) {}); + const int second = pool.Acquire([](int, int) {}); + + static_assert(noexcept(pool.Release(first))); + EXPECT_NO_THROW(pool.Release(first)); + EXPECT_NO_THROW(pool.Release(second)); + EXPECT_EQ(pool.ActiveCount(), 0u); + EXPECT_EQ(pool.FreeCount(), 2u); +} + +} // namespace +} // namespace Generators From 1085dcd4a0f13bb006f2080a3502570d45049a01 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Fri, 21 Aug 2026 18:23:25 -0500 Subject: [PATCH 13/19] Harden Engine request lifetime ownership Linearize external handle lifecycle state in the shared C API owner so Engine reclamation cannot race handle reacquisition, while replacing broad Request/Engine friendships with narrow internal capabilities. Files changed: - docs/paged_attention_engine.md - src/smartptrs.h - src/engine/engine.cpp - src/engine/engine.h - src/engine/request.cpp - src/engine/request.h - test/engine/engine_step_tests.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e54c6195-e3c4-4467-abcf-72bdff3260dc --- docs/paged_attention_engine.md | 6 +- src/engine/engine.cpp | 12 +-- src/engine/engine.h | 6 +- src/engine/request.cpp | 15 ++-- src/engine/request.h | 21 ++---- src/smartptrs.h | 80 +++++++++++++++----- test/engine/engine_step_tests.cpp | 119 ++++++++++++++++++++++++++++++ 7 files changed, 207 insertions(+), 52 deletions(-) diff --git a/docs/paged_attention_engine.md b/docs/paged_attention_engine.md index 6f1a12c9eb..605abb08a2 100644 --- a/docs/paged_attention_engine.md +++ b/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 status 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 status 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/`: @@ -157,6 +157,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. diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp index 4275e78277..a33ead3c25 100644 --- a/src/engine/engine.cpp +++ b/src/engine/engine.cpp @@ -75,7 +75,7 @@ void Engine::RemoveRequest(std::shared_ptr 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."); } @@ -87,7 +87,7 @@ void Engine::RemoveRequest(std::shared_ptr request) { ready_request_index_ = 0; std::erase(ready_requests_, request); std::erase(staged_ready_requests_, request); - request->CompleteClose(); + request->CompleteCloseFromEngine(*this); std::erase_if(tracked_requests_, [&request](const std::weak_ptr& tracked) { const auto owned = tracked.lock(); return !owned || owned == request; @@ -106,8 +106,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; @@ -115,7 +115,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); } } @@ -125,7 +125,7 @@ void Engine::ValidateRequestCanContinue(const std::shared_ptr& 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."); } diff --git a/src/engine/engine.h b/src/engine/engine.h index a8bca9718a..ab363952bc 100644 --- a/src/engine/engine.h +++ b/src/engine/engine.h @@ -118,12 +118,15 @@ struct Engine : std::enable_shared_from_this, */ bool HasPendingRequests() const; + // 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) const; + private: void ReclaimAbandonedRequests(); std::shared_ptr DrainReadyRequest(); std::shared_ptr StepDynamic(); std::shared_ptr StepStatic(); - void ValidateRequestCanContinue(const std::shared_ptr& request) const; [[noreturn]] void MarkUnhealthyAndThrow(StepOutcomeKind outcome, StepTransactionId transaction_id, const void* request_id, @@ -145,7 +148,6 @@ struct Engine : std::enable_shared_from_this, std::vector> staged_ready_requests_; size_t ready_request_index_{}; - friend struct Request; }; } // namespace Generators diff --git a/src/engine/request.cpp b/src/engine/request.cpp index d2c7e88e4c..5923c4a6b2 100644 --- a/src/engine/request.cpp +++ b/src/engine/request.cpp @@ -60,16 +60,13 @@ Request::Request(std::shared_ptr params) search_->DeferCompletion(true); } -void Request::OnFirstExternalReference() noexcept { - externally_abandoned_.store(false, std::memory_order_release); +bool Request::BelongsTo(const Engine& engine) const noexcept { + return engine_.lock().get() == &engine; } -void Request::OnLastExternalReference() noexcept { - externally_abandoned_.store(true, std::memory_order_release); -} - -bool Request::IsExternallyAbandoned() const noexcept { - return externally_abandoned_.load(std::memory_order_acquire); +void Request::CompleteCloseFromEngine(const Engine& engine) noexcept { + assert(BelongsTo(engine)); + CompleteClose(); } void Request::Assign(std::shared_ptr engine) { @@ -141,7 +138,7 @@ void Request::Remove() { engine->RemoveRequest(shared_from_this()); } -void Request::CompleteClose() { +void Request::CompleteClose() noexcept { engine_.reset(); status_ = RequestStatus::Closed; } diff --git a/src/engine/request.h b/src/engine/request.h index a81023af99..3f83f8a657 100644 --- a/src/engine/request.h +++ b/src/engine/request.h @@ -16,13 +16,6 @@ namespace Generators { -struct Request; - -template <> -struct ExternalRefCountedTraits { - static constexpr bool notify_external_reference_changes = true; -}; - struct RequestStepResult { int32_t token{}; bool token_appended{}; @@ -185,6 +178,11 @@ struct Request : std::enable_shared_from_this, */ void Remove(); + // Internal lifecycle capabilities used by Engine orchestration. These keep ownership and terminal + // mutation inside Request instead of exposing its weak owner or granting Engine private access. + bool BelongsTo(const Engine& engine) const noexcept; + void CompleteCloseFromEngine(const Engine& engine) noexcept; + /** * @brief Checks if the request is in prefill mode. * @return True while the tokens the application supplied have not all been through the model. @@ -324,13 +322,7 @@ struct Request : std::enable_shared_from_this, std::vector unseen_token_indices_; size_t next_unseen_token_index_{}; int64_t seen_sequence_length_{}; - friend struct Engine; - friend struct ExternalRefCounted; - - void CompleteClose(); - void OnFirstExternalReference() noexcept; - void OnLastExternalReference() noexcept; - bool IsExternallyAbandoned() const noexcept; + void CompleteClose() noexcept; int64_t processed_sequence_length_{}; // Sequence length the application's tokens reach up to. Everything below it is prompt, so the @@ -341,7 +333,6 @@ struct Request : std::enable_shared_from_this, std::unique_ptr search_; std::unique_ptr batched_sampler_state_; std::weak_ptr engine_; - std::atomic externally_abandoned_{false}; void ApplyLogitsProcessors(DeviceSpan logits); void SelectNextToken(); diff --git a/src/smartptrs.h b/src/smartptrs.h index a78ea947f8..2e57024c93 100644 --- a/src/smartptrs.h +++ b/src/smartptrs.h @@ -7,7 +7,9 @@ #include #include #include +#include #include // for std::remove_const_t +#include #include "span.h" #include "models/onnxruntime_api.h" // for ONNXTensorElementDataType #include "provider_options.h" // for ProviderOptions @@ -256,37 +258,77 @@ struct DeviceInterface { // A shared_ptr based type that we expose through our C API should inherit from this type. // ExternalAddRef must be called when returning an object through the C API -// ExternalRelease must be called on the C API destroy method -template -struct ExternalRefCountedTraits { - static constexpr bool notify_external_reference_changes = false; -}; - +// ExternalRelease must be called on the C API destroy method. template struct ExternalRefCounted { void ExternalAddRef() { - if (++ref_count_ == 1) { // First reference? - external_owner_ = static_cast(this)->shared_from_this(); - if constexpr (ExternalRefCountedTraits::notify_external_reference_changes) { - static_cast(this)->OnFirstExternalReference(); - } + ExternalReferenceLock lock{*this}; + if (ref_count_ == 0) { + // Acquire the self-owner before publishing the first reference. If shared_from_this throws, + // the never-acquired/zero-reference state remains unchanged. + auto owner = static_cast(this)->shared_from_this(); + external_owner_ = std::move(owner); + ref_count_ = 1; + external_lifecycle_started_ = true; + } else { + ++ref_count_; } } - void ExternalRelease() noexcept(ExternalRefCountedTraits::notify_external_reference_changes) { - if (--ref_count_ == 0) { - if constexpr (ExternalRefCountedTraits::notify_external_reference_changes) { - // Notify before releasing the self-owner so a type-specific last-release hook can only mark - // deferred work while the object is guaranteed to still be alive. - static_cast(this)->OnLastExternalReference(); + void ExternalRelease() noexcept { + std::shared_ptr released_owner; + { + ExternalReferenceLock lock{*this}; + assert(ref_count_ > 0); + if (--ref_count_ == 0) { + released_owner = std::move(external_owner_); } - external_owner_ = nullptr; } + // The self-owner may be the last strong reference. Destroy it only after releasing the member + // lock so object destruction never runs while code still accesses this object's synchronization. + } + + // True only after an external lifecycle has started and its final handle has been released. + // A never-exposed object therefore remains distinct from an abandoned external object. + bool ExternalReferencesAbandoned() const noexcept { + ExternalReferenceLock lock{*this}; + return external_lifecycle_started_ && ref_count_ == 0; } private: + void LockExternalReferences() const noexcept { + while (external_reference_lock_.test_and_set(std::memory_order_acquire)) { +#if defined(USE_CXX17) + std::this_thread::yield(); +#else + external_reference_lock_.wait(true, std::memory_order_relaxed); +#endif + } + } + + void UnlockExternalReferences() const noexcept { + external_reference_lock_.clear(std::memory_order_release); +#if !defined(USE_CXX17) + external_reference_lock_.notify_one(); +#endif + } + + struct ExternalReferenceLock { + explicit ExternalReferenceLock(const ExternalRefCounted& owner) noexcept + : owner_{owner} { + owner_.LockExternalReferences(); + } + ~ExternalReferenceLock() noexcept { + owner_.UnlockExternalReferences(); + } + + const ExternalRefCounted& owner_; + }; + std::shared_ptr external_owner_; // shared_ptr to ourselves to keep us alive - std::atomic ref_count_{}; // C API refcount (can't use only the shared_ptr) + int ref_count_{}; // Guarded with external_owner_ and lifecycle state. + bool external_lifecycle_started_{}; // Distinguishes never exposed from finally released. + mutable std::atomic_flag external_reference_lock_ = ATOMIC_FLAG_INIT; }; namespace Location { diff --git a/test/engine/engine_step_tests.cpp b/test/engine/engine_step_tests.cpp index 9ed65a6049..e70dc81f9a 100644 --- a/test/engine/engine_step_tests.cpp +++ b/test/engine/engine_step_tests.cpp @@ -9,7 +9,11 @@ // without redundant model runs, and forms a fresh batch across steps under capacity backpressure. #include +#include +#include #include +#include +#include #include #include @@ -19,6 +23,41 @@ #include "engine_test_doubles.h" namespace Generators { + +class TestBarrier { + public: + explicit TestBarrier(size_t participant_count) + : remaining_{participant_count} {} + + void ArriveAndWait() { + std::unique_lock lock{mutex_}; + if (--remaining_ == 0) { + condition_.notify_all(); + return; + } + condition_.wait(lock, [this] { return remaining_ == 0; }); + } + + private: + std::mutex mutex_; + std::condition_variable condition_; + size_t remaining_; +}; + +struct ExternalRequestRaceProbe + : std::enable_shared_from_this, + ExternalRefCounted { + explicit ExternalRequestRaceProbe( + std::shared_ptr> destroyed) + : destroyed_{std::move(destroyed)} {} + + ~ExternalRequestRaceProbe() { + destroyed_->store(true, std::memory_order_release); + } + + std::shared_ptr> destroyed_; +}; + namespace test { namespace { @@ -58,6 +97,8 @@ class ExternalRequestReference { }; static_assert(noexcept(std::declval().ExternalRelease())); +static_assert( + noexcept(std::declval().ExternalRelease())); class EngineStepTest : public ::testing::Test { protected: @@ -70,6 +111,52 @@ class EngineStepTest : public ::testing::Test { std::shared_ptr model_; }; +TEST(ExternalRefCountedTest, + DistinguishesNeverHeldHeldAbandonedAndReacquiredStates) { + auto destroyed = std::make_shared>(false); + auto probe = std::make_shared(destroyed); + + EXPECT_FALSE(probe->ExternalReferencesAbandoned()); + probe->ExternalAddRef(); + EXPECT_FALSE(probe->ExternalReferencesAbandoned()); + probe->ExternalRelease(); + EXPECT_TRUE(probe->ExternalReferencesAbandoned()); + probe->ExternalAddRef(); + EXPECT_FALSE(probe->ExternalReferencesAbandoned()); + probe->ExternalRelease(); + EXPECT_TRUE(probe->ExternalReferencesAbandoned()); +} + +TEST(ExternalRefCountedTest, + ConcurrentFinalReleaseAndReacquirePreserveOwnerLifetime) { + auto destroyed = std::make_shared>(false); + auto engine_owner = + std::make_shared(destroyed); + auto* raw = engine_owner.get(); + raw->ExternalAddRef(); + + TestBarrier transition_start{3}; + std::thread release_thread([raw, &transition_start] { + transition_start.ArriveAndWait(); + raw->ExternalRelease(); + }); + std::thread reacquire_thread([engine_owner, &transition_start] { + transition_start.ArriveAndWait(); + engine_owner->ExternalAddRef(); + }); + transition_start.ArriveAndWait(); + release_thread.join(); + reacquire_thread.join(); + + EXPECT_FALSE(raw->ExternalReferencesAbandoned()); + EXPECT_FALSE(destroyed->load(std::memory_order_acquire)); + + engine_owner.reset(); + EXPECT_FALSE(destroyed->load(std::memory_order_acquire)); + raw->ExternalRelease(); + EXPECT_TRUE(destroyed->load(std::memory_order_acquire)); +} + // One request: Step decodes the proposed batch exactly once, commits its cache allocation, and // returns the request. TEST_F(EngineStepTest, SingleRequestSchedulesThenDecodesThenReturns) { @@ -251,6 +338,38 @@ TEST_F(EngineStepTest, ReacquiringExternalReferenceCancelsDeferredAbandonment) { reacquired_external.Release(); } +TEST_F(EngineStepTest, + ConcurrentFinalReleaseAndReacquireCancelsRequestAbandonment) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/1, EosToken(*model_)); + auto prompt = Prompt(10); + auto request = MintRequest(*model_, prompt); + request->ExternalAddRef(); + engine.engine->AddRequest(request); + ASSERT_EQ(engine.engine->Step(), request); + ASSERT_EQ(request->status_, RequestStatus::TurnComplete); + + TestBarrier transition_start{3}; + std::thread release_thread([request, &transition_start] { + transition_start.ArriveAndWait(); + request->ExternalRelease(); + }); + std::thread reacquire_thread([request, &transition_start] { + transition_start.ArriveAndWait(); + request->ExternalAddRef(); + }); + transition_start.ArriveAndWait(); + release_thread.join(); + reacquire_thread.join(); + + EXPECT_EQ(engine.engine->Step(), nullptr); + EXPECT_EQ(request->status_, RequestStatus::TurnComplete); + EXPECT_EQ(engine.cache->AllocatedCount(), 1u); + EXPECT_EQ(engine.cache->deallocate_calls, 0); + + engine.engine->RemoveRequest(request); + request->ExternalRelease(); +} + TEST_F(EngineStepTest, ContinueRejectsUndrainedReadyNotificationWithoutMutation) { auto engine = MakeDoublesEngine(model_, /*capacity=*/8, EosToken(*model_)); auto first_prompt = Prompt(10); From 0248bb42eb4f7653c887c1516d2a602d6060ef3a Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Fri, 21 Aug 2026 18:23:44 -0500 Subject: [PATCH 14/19] Release sampler indices after allocation failure Make CUDA sampler index acquisition and wrapper construction one transaction so a post-acquire allocation failure returns the index without leaking pool capacity. Files changed: - src/cuda/interface.cpp - src/cuda/sampler_state_index_pool.h - test/engine/sampler_state_index_pool_tests.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e54c6195-e3c4-4467-abcf-72bdff3260dc --- src/cuda/interface.cpp | 30 ++++++++++++------- src/cuda/sampler_state_index_pool.h | 11 +++++++ .../engine/sampler_state_index_pool_tests.cpp | 26 ++++++++++++++++ 3 files changed, 56 insertions(+), 11 deletions(-) diff --git a/src/cuda/interface.cpp b/src/cuda/interface.cpp index efb61ca714..72b141b6af 100644 --- a/src/cuda/interface.cpp +++ b/src/cuda/interface.cpp @@ -93,16 +93,19 @@ struct CudaSamplerStatePool { } } - int Acquire(int random_seed) { - return indices_.Acquire([this, random_seed](int index, int required_size) { - EnsureCapacity(required_size); - const unsigned long long seed = - random_seed == -1 - ? static_cast(std::random_device{}()) - : static_cast(random_seed); - cuda::LaunchInitCurandState( - seed, states_.Span().data() + index, GetStream()); - }); + template + 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(std::random_device{}()) + : static_cast(random_seed); + cuda::LaunchInitCurandState( + seed, states_.Span().data() + index, GetStream()); + }, + std::forward(create)); } void Release(int index) noexcept { indices_.Release(index); } @@ -157,7 +160,12 @@ struct CudaBatchedSampler final : BatchedSampler { } std::unique_ptr CreateState(int random_seed) override { - return std::make_unique(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(pool, index); + }); } bool OwnsState(const BatchedSamplerState& state) const override { diff --git a/src/cuda/sampler_state_index_pool.h b/src/cuda/sampler_state_index_pool.h index aa3b429255..e6402be4ac 100644 --- a/src/cuda/sampler_state_index_pool.h +++ b/src/cuda/sampler_state_index_pool.h @@ -32,6 +32,17 @@ class SamplerStateIndexPool { return index; } + template + auto AcquireOwned(Prepare&& prepare, Create&& create) { + const int index = Acquire(std::forward(prepare)); + try { + return std::forward(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) == diff --git a/test/engine/sampler_state_index_pool_tests.cpp b/test/engine/sampler_state_index_pool_tests.cpp index 6bacda9eb2..b71b8bfb4c 100644 --- a/test/engine/sampler_state_index_pool_tests.cpp +++ b/test/engine/sampler_state_index_pool_tests.cpp @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include @@ -51,6 +52,31 @@ TEST(SamplerStateIndexPoolTest, ReusedIndexPreparationFailureKeepsItFree) { EXPECT_EQ(pool.ActiveCount(), 1u); } +TEST(SamplerStateIndexPoolTest, OwnedStateConstructionFailureReleasesAcquiredIndex) { + SamplerStateIndexPool pool; + + EXPECT_THROW( + pool.AcquireOwned( + [](int, int) {}, + [](int) -> std::unique_ptr { + throw std::bad_alloc{}; + }), + std::bad_alloc); + + EXPECT_EQ(pool.Size(), 1); + EXPECT_EQ(pool.FreeCount(), 1u); + EXPECT_EQ(pool.ActiveCount(), 0u); + + auto state = pool.AcquireOwned( + [](int, int) {}, + [](int index) { + return std::make_unique(index); + }); + ASSERT_NE(state, nullptr); + EXPECT_EQ(*state, 0); + EXPECT_EQ(pool.ActiveCount(), 1u); +} + TEST(SamplerStateIndexPoolTest, ReleaseIsNoexceptAndAllocationFree) { SamplerStateIndexPool pool; const int first = pool.Acquire([](int, int) {}); From bda27927807c70c56dea368b4521b7e3c7f818ae Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Fri, 21 Aug 2026 18:47:59 -0500 Subject: [PATCH 15/19] Adapt hardening to current Engine recovery Preserve friend-free continuation rollback handling and restore diagnostic snapshot fixtures after merging the latest main recovery and unseen-output changes. Files changed: - src/engine/engine.cpp - src/engine/engine.h - test/engine/engine_invariants_tests.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e54c6195-e3c4-4467-abcf-72bdff3260dc --- src/engine/engine.cpp | 2 +- src/engine/engine.h | 8 +- test/engine/engine_invariants_tests.cpp | 146 +++++++++++++++++++++++- 3 files changed, 147 insertions(+), 9 deletions(-) diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp index 41abfe3bf2..e9dc7164c1 100644 --- a/src/engine/engine.cpp +++ b/src/engine/engine.cpp @@ -172,7 +172,7 @@ void Engine::ValidateRequestCanContinue(const std::shared_ptr& request) message = AddExceptionCause( std::move(message) + " Closing the poisoned request also failed.", std::current_exception()); - request->CompleteClose(); + request->CompleteCloseFromEngine(*this); } MarkUnhealthyAndThrow( StepOutcomeKind::FatalExecutionFailure, diff --git a/src/engine/engine.h b/src/engine/engine.h index 804818892b..774ffbd7fe 100644 --- a/src/engine/engine.h +++ b/src/engine/engine.h @@ -121,16 +121,16 @@ struct Engine : std::enable_shared_from_this, // 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) const; + [[noreturn]] void HandleContinuationRestoreFailure( + const std::shared_ptr& request, + std::exception_ptr append_error, + std::exception_ptr restore_error); private: void ReclaimAbandonedRequests(); std::shared_ptr DrainReadyRequest(); std::shared_ptr StepDynamic(); std::shared_ptr StepStatic(); - [[noreturn]] void HandleContinuationRestoreFailure( - const std::shared_ptr& request, - std::exception_ptr append_error, - std::exception_ptr restore_error); [[noreturn]] void MarkUnhealthyAndThrow(StepOutcomeKind outcome, StepTransactionId transaction_id, const void* request_id, diff --git a/test/engine/engine_invariants_tests.cpp b/test/engine/engine_invariants_tests.cpp index 5e9ff01650..4c799b6e46 100644 --- a/test/engine/engine_invariants_tests.cpp +++ b/test/engine/engine_invariants_tests.cpp @@ -102,7 +102,14 @@ TEST(InvariantValidatorTest, WindowBlockPoolValidatesRingSizeAndOwnershipIndepen } TEST(InvariantValidatorTest, WindowBlockPoolIncludesTransactionReservationsInAccounting) { - auto cache = MakeValidCache(); + PagedCacheSnapshot cache; + cache.block_size = kBlockSize; + cache.total_blocks = 1; + cache.transaction_reserved_block_ids = {0}; + cache.reservations = { + RequestReservationSnapshot{ + kRequestA, 0, 1, 0, {0}, {0, 1}, true}, + }; cache.window_blocks.total_blocks = 2; cache.window_blocks.free_blocks = 0; cache.window_blocks.blocks_per_request = 2; @@ -148,7 +155,7 @@ TEST(InvariantValidatorTest, InitialAdmissionReservationValidatesWithoutCommitte cache.free_blocks = 0; cache.transaction_reserved_block_ids = {0}; cache.reservations = { - RequestReservationSnapshot{kRequestA, 0, 1, 0, {0}}, + RequestReservationSnapshot{kRequestA, 0, 1, 0, {0}, {}, true}, }; EXPECT_TRUE(ValidateCacheInvariants(cache).empty()); @@ -176,7 +183,7 @@ TEST(InvariantValidatorTest, InitialAdmissionRejectsUnreservedDeltaBlock) { cache.free_blocks = 1; cache.transaction_reserved_block_ids = {0}; cache.reservations = { - RequestReservationSnapshot{kRequestA, 0, 1, 0, {1}}, + RequestReservationSnapshot{kRequestA, 0, 1, 0, {1}, {}, true}, }; const auto violations = ValidateCacheInvariants(cache); @@ -329,7 +336,7 @@ TEST(InvariantValidatorTest, TurnCompleteRequestFullyProcessedIsValid) { TEST(InvariantValidatorTest, ConsistentSnapshotsValidateClean) { const auto cache = MakeValidCache(); const std::vector requests{ - MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 5, 9), MakeValidRequest(kRequestB, RequestStatus::Active, 4, 4, 4), }; EXPECT_TRUE(ValidateInvariants(cache, requests).empty()); @@ -345,6 +352,137 @@ TEST(InvariantValidatorTest, BlockTableForUnknownRequestReported) { EXPECT_FALSE(ValidateInvariants(cache, requests).empty()); } +TEST(InvariantValidatorTest, CommittedCacheUsageMustMatchProcessedLength) { + const auto cache = MakeValidCache(); + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 8, 8), + MakeValidRequest(kRequestB, RequestStatus::Active, 4, 4, 4), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "differs from processed sequence length") != + std::string::npos; + }), + violations.end()); +} + +TEST(InvariantValidatorTest, FullAndWindowOwnerSetsMustAgree) { + auto cache = MakeValidCache(); + cache.window_blocks.total_blocks = 2; + cache.window_blocks.blocks_per_request = 2; + cache.window_blocks.requests = { + RequestBlockSnapshot{kRequestA, {0, 1}}, + }; + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + MakeValidRequest(kRequestB, RequestStatus::Active, 4, 4, 4), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "Full-cache and window-cache owner sets disagree") != + std::string::npos; + }), + violations.end()); +} + +TEST(InvariantValidatorTest, WindowOwnerMustBeAKnownRequest) { + auto cache = MakeValidCache(); + cache.window_blocks.total_blocks = 4; + cache.window_blocks.blocks_per_request = 2; + cache.window_blocks.requests = { + RequestBlockSnapshot{kRequestA, {0, 1}}, + RequestBlockSnapshot{kRequestB, {2, 3}}, + }; + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "Window cache holds a block table for unknown Request") != + std::string::npos; + }), + violations.end()); +} + +TEST(InvariantValidatorTest, ReservationOwnerMustBeKnownAndMatchCommittedUsage) { + auto cache = MakeValidCache(); + cache.free_blocks = 0; + cache.transaction_reserved_block_ids = {3}; + cache.reservations = { + RequestReservationSnapshot{ + kRequestB, /*committed_slots=*/3, /*target_slots=*/5, + /*tail_slots_to_consume=*/0, /*reserved_block_ids=*/{3}}, + }; + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Active, 9, 9, 9), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "transaction committed slots disagree") != + std::string::npos; + }), + violations.end()); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "transaction reservation for unknown Request") != + std::string::npos; + }), + violations.end()); +} + +TEST(InvariantValidatorTest, WindowReservationMustHaveConsistentOwnership) { + PagedCacheSnapshot cache; + cache.block_size = kBlockSize; + cache.total_blocks = 1; + cache.transaction_reserved_block_ids = {0}; + cache.window_blocks.total_blocks = 2; + cache.window_blocks.blocks_per_request = 2; + cache.window_blocks.transaction_reserved_block_ids = {0, 1}; + cache.reservations = { + RequestReservationSnapshot{ + kRequestA, 0, 1, 0, {0}, {0}, true}, + }; + const std::vector requests{ + MakeValidRequest(kRequestA, RequestStatus::Assigned, 3, 0, 3), + }; + + const auto violations = ValidateInvariants(cache, requests); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "Not every transaction-reserved window block belongs") != + std::string::npos; + }), + violations.end()); + EXPECT_NE(std::find_if( + violations.begin(), violations.end(), + [](const InvariantViolation& violation) { + return violation.message.find( + "transaction-reserved window blocks instead of 2") != + std::string::npos; + }), + violations.end()); +} + TEST(InvariantValidatorTest, ThrowWrapperListsViolations) { auto cache = MakeValidCache(); cache.free_blocks = 0; // break block accounting From a66b59821c2c2fce6dc41c16d3f0acc37c6eac4c Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Fri, 21 Aug 2026 19:02:37 -0500 Subject: [PATCH 16/19] Use stable Python turn completion API Keep executable windowed continuation coverage on current main without exposing queued scheduler status through the Python Request surface. Files changed: - test/python/models/test_engine_windowed_multiwrap.py Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e54c6195-e3c4-4467-abcf-72bdff3260dc --- test/python/models/test_engine_windowed_multiwrap.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/python/models/test_engine_windowed_multiwrap.py b/test/python/models/test_engine_windowed_multiwrap.py index 454f01ac1b..dda907a08f 100644 --- a/test/python/models/test_engine_windowed_multiwrap.py +++ b/test/python/models/test_engine_windowed_multiwrap.py @@ -162,7 +162,7 @@ def test_continuation_crosses_two_window_ring_wraps_and_matches_clean_replay(tes first_turn = _run_turn(engine, request) assert first_turn == _EXPECTED_FIRST_TURN - assert request.status == og.RequestStatus.TURN_COMPLETE + assert request.is_turn_complete() first_turn_length = len(_INITIAL_PROMPT) + len(first_turn) assert first_turn_length == 3 assert first_turn_length < _MAX_LENGTH @@ -172,12 +172,12 @@ def test_continuation_crosses_two_window_ring_wraps_and_matches_clean_replay(tes # the full/windowed caches disagree. Tokens at positions 11 and 12 check # both columns of the two-block ring before EOS at position 13. request.continue_with(np.asarray(_CONTINUATION, dtype=np.int32)) - assert request.status == og.RequestStatus.QUEUED + assert not request.is_turn_complete() second_turn = _run_turn(engine, request) assert second_turn == _EXPECTED_SECOND_TURN assert _INVARIANT_FAILURE_TOKEN_ID not in second_turn - assert request.status == og.RequestStatus.TURN_COMPLETE + assert request.is_turn_complete() retained_length = first_turn_length + len(_CONTINUATION) + len(second_turn) assert retained_length == 14 assert retained_length < _MAX_LENGTH From c8b57d4e7e9439b2a92215a0f1420f55a0c1f857 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Mon, 24 Aug 2026 16:24:25 -0500 Subject: [PATCH 17/19] Format Engine hardening sources Apply the repository-pinned clang-format 20.1.0 rules to the PR files so repository-wide C++ format lint passes. Files changed: - src/engine/engine.h - src/engine/paged_key_value_cache.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e54c6195-e3c4-4467-abcf-72bdff3260dc --- src/engine/engine.h | 1 - src/engine/paged_key_value_cache.cpp | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/src/engine/engine.h b/src/engine/engine.h index 774ffbd7fe..44b9067db4 100644 --- a/src/engine/engine.h +++ b/src/engine/engine.h @@ -151,7 +151,6 @@ struct Engine : std::enable_shared_from_this, std::vector> ready_requests_; std::vector> staged_ready_requests_; size_t ready_request_index_{}; - }; } // namespace Generators diff --git a/src/engine/paged_key_value_cache.cpp b/src/engine/paged_key_value_cache.cpp index 3f87f4b8ad..f281cef790 100644 --- a/src/engine/paged_key_value_cache.cpp +++ b/src/engine/paged_key_value_cache.cpp @@ -593,8 +593,7 @@ PagedCacheSnapshot PagedKeyValueCache::Snapshot( delta.reserved_window_block_count); for (size_t i = 0; i < delta.reserved_window_block_count; ++i) { request_reservation.reserved_window_block_ids.push_back( - reservation.ReservedWindowBlocks()[ - delta.reserved_window_block_offset + i] + reservation.ReservedWindowBlocks()[delta.reserved_window_block_offset + i] ->Id()); } request_reservation.newly_admitted = delta.newly_admitted; From 232d000d3e85896e149db86e86262a5da5e6d10f Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Mon, 24 Aug 2026 17:49:10 -0500 Subject: [PATCH 18/19] Inline CUDA sampler index management Keep failure-atomic sampler index ownership local to the CUDA implementation while preserving allocation-free rollback and removing the dedicated helper surface and tests. Files changed: - src/cuda/interface.cpp - src/cuda/sampler_state_index_pool.h - test/engine/sampler_state_index_pool_tests.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e54c6195-e3c4-4467-abcf-72bdff3260dc --- src/cuda/interface.cpp | 52 ++++++++++- src/cuda/sampler_state_index_pool.h | 65 ------------- .../engine/sampler_state_index_pool_tests.cpp | 93 ------------------- 3 files changed, 51 insertions(+), 159 deletions(-) delete mode 100644 src/cuda/sampler_state_index_pool.h delete mode 100644 test/engine/sampler_state_index_pool_tests.cpp diff --git a/src/cuda/interface.cpp b/src/cuda/interface.cpp index 26bd23eebe..a0910ca6fa 100644 --- a/src/cuda/interface.cpp +++ b/src/cuda/interface.cpp @@ -8,13 +8,17 @@ #include "search_cuda.h" #include "kernels.h" #include "cuda_topk.h" -#include "sampler_state_index_pool.h" +#include +#include #include +#include #include #include #include #include #include +#include +#include #if defined(_WIN32) || defined(_WIN64) #define strcasecmp _stricmp @@ -90,6 +94,52 @@ DeviceSpan AllocateCudaSpan(size_t count) { return DeviceSpan{std::make_shared(count * sizeof(T))}; } +namespace { + +class SamplerStateIndexPool { + public: + template + 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(required_size)); + std::forward(prepare)(index, required_size); + + if (reusing) { + free_indices_.pop_back(); + } else { + size_ = required_size; + } + + try { + return std::forward(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 free_indices_; + int size_{}; +}; + +} // namespace + struct CudaSamplerStatePool { explicit CudaSamplerStatePool(int initial_capacity) { if (initial_capacity > 0) { diff --git a/src/cuda/sampler_state_index_pool.h b/src/cuda/sampler_state_index_pool.h deleted file mode 100644 index e6402be4ac..0000000000 --- a/src/cuda/sampler_state_index_pool.h +++ /dev/null @@ -1,65 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include -#include -#include -#include -#include - -namespace Generators { - -class SamplerStateIndexPool { - public: - template - int Acquire(Prepare&& prepare) { - 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 noexcept, so reserve its future slot before any external preparation can - // publish an acquired index. A throwing preparation leaves every pool counter unchanged. - free_indices_.reserve(static_cast(required_size)); - std::forward(prepare)(index, required_size); - - if (reusing) { - free_indices_.pop_back(); - } else { - size_ = required_size; - } - return index; - } - - template - auto AcquireOwned(Prepare&& prepare, Create&& create) { - const int index = Acquire(std::forward(prepare)); - try { - return std::forward(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_; } - size_t FreeCount() const noexcept { return free_indices_.size(); } - size_t ActiveCount() const noexcept { - return static_cast(size_) - free_indices_.size(); - } - - private: - std::vector free_indices_; - int size_{}; -}; - -} // namespace Generators diff --git a/test/engine/sampler_state_index_pool_tests.cpp b/test/engine/sampler_state_index_pool_tests.cpp deleted file mode 100644 index b71b8bfb4c..0000000000 --- a/test/engine/sampler_state_index_pool_tests.cpp +++ /dev/null @@ -1,93 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include -#include -#include - -#include - -#include "cuda/sampler_state_index_pool.h" - -namespace Generators { -namespace { - -TEST(SamplerStateIndexPoolTest, NewIndexPreparationFailureIsAtomic) { - SamplerStateIndexPool pool; - - EXPECT_THROW( - pool.Acquire([](int index, int required_size) { - EXPECT_EQ(index, 0); - EXPECT_EQ(required_size, 1); - throw std::runtime_error("Injected sampler state growth failure."); - }), - std::runtime_error); - EXPECT_EQ(pool.Size(), 0); - EXPECT_EQ(pool.ActiveCount(), 0u); - EXPECT_EQ(pool.FreeCount(), 0u); - - const int index = - pool.Acquire([](int, int) {}); - EXPECT_EQ(index, 0); - EXPECT_EQ(pool.Size(), 1); - EXPECT_EQ(pool.ActiveCount(), 1u); -} - -TEST(SamplerStateIndexPoolTest, ReusedIndexPreparationFailureKeepsItFree) { - SamplerStateIndexPool pool; - const int original = pool.Acquire([](int, int) {}); - pool.Release(original); - - EXPECT_THROW( - pool.Acquire([](int, int) { - throw std::runtime_error("Injected sampler state initialization failure."); - }), - std::runtime_error); - EXPECT_EQ(pool.Size(), 1); - EXPECT_EQ(pool.ActiveCount(), 0u); - EXPECT_EQ(pool.FreeCount(), 1u); - - const int retried = pool.Acquire([](int, int) {}); - EXPECT_EQ(retried, original); - EXPECT_EQ(pool.ActiveCount(), 1u); -} - -TEST(SamplerStateIndexPoolTest, OwnedStateConstructionFailureReleasesAcquiredIndex) { - SamplerStateIndexPool pool; - - EXPECT_THROW( - pool.AcquireOwned( - [](int, int) {}, - [](int) -> std::unique_ptr { - throw std::bad_alloc{}; - }), - std::bad_alloc); - - EXPECT_EQ(pool.Size(), 1); - EXPECT_EQ(pool.FreeCount(), 1u); - EXPECT_EQ(pool.ActiveCount(), 0u); - - auto state = pool.AcquireOwned( - [](int, int) {}, - [](int index) { - return std::make_unique(index); - }); - ASSERT_NE(state, nullptr); - EXPECT_EQ(*state, 0); - EXPECT_EQ(pool.ActiveCount(), 1u); -} - -TEST(SamplerStateIndexPoolTest, ReleaseIsNoexceptAndAllocationFree) { - SamplerStateIndexPool pool; - const int first = pool.Acquire([](int, int) {}); - const int second = pool.Acquire([](int, int) {}); - - static_assert(noexcept(pool.Release(first))); - EXPECT_NO_THROW(pool.Release(first)); - EXPECT_NO_THROW(pool.Release(second)); - EXPECT_EQ(pool.ActiveCount(), 0u); - EXPECT_EQ(pool.FreeCount(), 2u); -} - -} // namespace -} // namespace Generators From b927e4ce2010ffad40ad0028109ea4d053f480d9 Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Tue, 25 Aug 2026 19:44:28 -0500 Subject: [PATCH 19/19] Remove windowed cache diagnostic/coverage scope creep Splits the paged/full/window cache diagnostic invariant strengthening and the new synthetic-windowed-multiwrap ring-write/rollback-retry test coverage out of this PR. That work is a separate cache-correctness concern, not required for the concurrent final-release/reacquisition race fix or sampler admission atomicity this PR targets. It now lives on the bmehta001-windowed-cache-diagnostics branch for a follow-up PR. Files reverted to main / removed: - .gitignore - src/engine/engine_invariants.cpp/.h - src/engine/paged_key_value_cache.cpp - test/engine/windowed_transaction_tests.cpp - test/models/engine/synthetic-windowed-multiwrap/* - test/python/create/create_synthetic_windowed_multiwrap_model.py - test/python/models/test_engine_windowed_multiwrap.py Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .gitignore | 6 - src/engine/engine_invariants.cpp | 88 --- src/engine/engine_invariants.h | 2 - src/engine/paged_key_value_cache.cpp | 31 +- test/engine/windowed_transaction_tests.cpp | 514 -------------- .../synthetic-windowed-multiwrap/decoder.onnx | Bin 10403 -> 0 bytes .../genai_config.json | 56 -- ...eate_synthetic_windowed_multiwrap_model.py | 631 ------------------ .../models/test_engine_windowed_multiwrap.py | 199 ------ 9 files changed, 8 insertions(+), 1519 deletions(-) delete mode 100644 test/engine/windowed_transaction_tests.cpp delete mode 100644 test/models/engine/synthetic-windowed-multiwrap/decoder.onnx delete mode 100644 test/models/engine/synthetic-windowed-multiwrap/genai_config.json delete mode 100644 test/python/create/create_synthetic_windowed_multiwrap_model.py delete mode 100644 test/python/models/test_engine_windowed_multiwrap.py diff --git a/.gitignore b/.gitignore index bc70ee352d..ec7a8ed488 100644 --- a/.gitignore +++ b/.gitignore @@ -44,12 +44,6 @@ examples/csharp/ModelChat/models !test/models/qwen3-5/* !test/models/qwen3-vl/* !test/models/whisper/* -!/test/models/engine/ -/test/models/engine/* -!/test/models/engine/synthetic-windowed-multiwrap/ -/test/models/engine/synthetic-windowed-multiwrap/* -!/test/models/engine/synthetic-windowed-multiwrap/decoder.onnx -!/test/models/engine/synthetic-windowed-multiwrap/genai_config.json .ipynb_checkpoints/ /src/java/.gradle diff --git a/src/engine/engine_invariants.cpp b/src/engine/engine_invariants.cpp index adeaf121f4..10552265a9 100644 --- a/src/engine/engine_invariants.cpp +++ b/src/engine/engine_invariants.cpp @@ -115,7 +115,6 @@ std::vector ValidateCacheInvariants(const PagedCacheSnapshot std::unordered_set reservation_requests; std::unordered_set blocks_assigned_to_delta; - std::unordered_set window_blocks_assigned_to_delta; for (const auto& reservation : cache.reservations) { if (!reservation_requests.insert(reservation.request_id).second) { add("Request " + PtrId(reservation.request_id) + @@ -130,20 +129,6 @@ std::vector ValidateCacheInvariants(const PagedCacheSnapshot add("Request " + PtrId(reservation.request_id) + " transaction tail-slot growth exceeds total growth."); } - const auto committed_owner = std::find_if( - cache.requests.begin(), cache.requests.end(), - [&reservation](const RequestBlockSnapshot& request) { - return request.request_id == reservation.request_id; - }); - if (reservation.newly_admitted == (committed_owner != cache.requests.end())) { - add("Request " + PtrId(reservation.request_id) + - " transaction membership disagrees with committed cache ownership."); - } - if (committed_owner != cache.requests.end() && - reservation.committed_slots != committed_owner->used_slots) { - add("Request " + PtrId(reservation.request_id) + - " transaction committed slots disagree with committed cache usage."); - } for (const size_t block_id : reservation.reserved_block_ids) { if (reserved_blocks.find(block_id) == reserved_blocks.end()) { add("Request " + PtrId(reservation.request_id) + @@ -155,12 +140,6 @@ std::vector ValidateCacheInvariants(const PagedCacheSnapshot " is assigned to more than one Request delta."); } } - for (const size_t block_id : reservation.reserved_window_block_ids) { - if (!window_blocks_assigned_to_delta.insert(block_id).second) { - add("Transaction-reserved window block id " + std::to_string(block_id) + - " is assigned to more than one Request delta."); - } - } } if (blocks_assigned_to_delta != reserved_blocks) { add("Not every transaction-reserved block belongs to exactly one Request delta."); @@ -208,25 +187,6 @@ std::vector ValidateCacheInvariants(const PagedCacheSnapshot " is also committed to a Request."); } } - if (window_blocks_assigned_to_delta != reserved_window_blocks) { - add("Not every transaction-reserved window block belongs to exactly one Request delta."); - } - if (window.total_blocks == 0 && !window_blocks_assigned_to_delta.empty()) { - add("A non-windowed cache has Request-attributed window reservations."); - } - if (window.total_blocks != 0) { - for (const auto& reservation : cache.reservations) { - const size_t expected_window_blocks = - reservation.newly_admitted ? window.blocks_per_request : 0; - if (reservation.reserved_window_block_ids.size() != - expected_window_blocks) { - add("Request " + PtrId(reservation.request_id) + " owns " + - std::to_string(reservation.reserved_window_block_ids.size()) + - " transaction-reserved window blocks instead of " + - std::to_string(expected_window_blocks) + "."); - } - } - } if (window.free_blocks > window.total_blocks) { add("window free_blocks (" + std::to_string(window.free_blocks) + ") exceeds total_blocks (" + std::to_string(window.total_blocks) + ")."); @@ -295,54 +255,6 @@ std::vector ValidateInvariants(const PagedCacheSnapshot& cac "Cache holds a block table for unknown Request " + PtrId(owner.request_id) + "."}); } } - for (const auto& owner : cache.window_blocks.requests) { - if (known_requests.find(owner.request_id) == known_requests.end()) { - violations.push_back(InvariantViolation{ - "Window cache holds a block table for unknown Request " + - PtrId(owner.request_id) + "."}); - } - } - for (const auto& reservation : cache.reservations) { - if (known_requests.find(reservation.request_id) == known_requests.end()) { - violations.push_back(InvariantViolation{ - "Cache holds a transaction reservation for unknown Request " + - PtrId(reservation.request_id) + "."}); - } - } - - std::unordered_map processed_by_request; - for (const auto& request : requests) { - if (request.processed_sequence_length >= 0) { - processed_by_request.emplace( - request.request_id, - static_cast(request.processed_sequence_length)); - } - } - for (const auto& owner : cache.requests) { - const auto processed = processed_by_request.find(owner.request_id); - if (processed != processed_by_request.end() && - owner.used_slots != processed->second) { - violations.push_back(InvariantViolation{ - "Request " + PtrId(owner.request_id) + " committed cache usage (" + - std::to_string(owner.used_slots) + ") differs from processed sequence length (" + - std::to_string(processed->second) + ")."}); - } - } - - if (cache.window_blocks.total_blocks != 0) { - std::set full_owners; - std::set window_owners; - for (const auto& owner : cache.requests) { - full_owners.insert(owner.request_id); - } - for (const auto& owner : cache.window_blocks.requests) { - window_owners.insert(owner.request_id); - } - if (full_owners != window_owners) { - violations.push_back(InvariantViolation{ - "Full-cache and window-cache owner sets disagree."}); - } - } return violations; } diff --git a/src/engine/engine_invariants.h b/src/engine/engine_invariants.h index d69ed97294..164c56ed81 100644 --- a/src/engine/engine_invariants.h +++ b/src/engine/engine_invariants.h @@ -53,8 +53,6 @@ struct RequestReservationSnapshot { size_t target_slots{}; size_t tail_slots_to_consume{}; std::vector reserved_block_ids; - std::vector reserved_window_block_ids; - bool newly_admitted{}; }; struct WindowBlockPoolSnapshot { diff --git a/src/engine/paged_key_value_cache.cpp b/src/engine/paged_key_value_cache.cpp index f281cef790..acc34564f7 100644 --- a/src/engine/paged_key_value_cache.cpp +++ b/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,11 +46,9 @@ 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; } @@ -64,7 +56,8 @@ StateGroup ResolvePagedKeyValueGroup(const Config::Model::Decoder& decoder) { ONNXTensorElementDataType KeyValueCacheType(const std::shared_ptr& 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) 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(block_size, num_blocks); if (Windowed()) { @@ -589,14 +582,6 @@ PagedCacheSnapshot PagedKeyValueCache::Snapshot( request_reservation.reserved_block_ids.push_back( reservation.ReservedBlocks()[delta.reserved_block_offset + i]->Id()); } - request_reservation.reserved_window_block_ids.reserve( - delta.reserved_window_block_count); - for (size_t i = 0; i < delta.reserved_window_block_count; ++i) { - request_reservation.reserved_window_block_ids.push_back( - reservation.ReservedWindowBlocks()[delta.reserved_window_block_offset + i] - ->Id()); - } - request_reservation.newly_admitted = delta.newly_admitted; snapshot.reservations.push_back(std::move(request_reservation)); } return snapshot; diff --git a/test/engine/windowed_transaction_tests.cpp b/test/engine/windowed_transaction_tests.cpp deleted file mode 100644 index 3b88a34558..0000000000 --- a/test/engine/windowed_transaction_tests.cpp +++ /dev/null @@ -1,514 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include -#include -#include -#include -#include -#include -#include - -#include - -#include "engine/cache_manager.h" -#include "engine/engine.h" -#include "engine/engine_invariants.h" -#include "engine/model_executor.h" -#include "engine/scheduler.h" -#include "engine_test_helpers.h" - -namespace Generators { -namespace test { -namespace { - -constexpr int32_t kInvariantFailureToken = 31; -const std::vector kInitialPrompt{4, 6}; -const std::vector kFirstTurnOutput{12}; -const std::vector kContinuation{8, 3, 10, 5, 12, 7, 14, 9, 16}; -const std::vector kSecondTurnOutput{28, 17}; - -std::shared_ptr LoadWindowedMultiwrapModel() { - return CreateModel( - GetOrtEnv(), MODEL_PATH "engine/synthetic-windowed-multiwrap"); -} - -std::vector RunTurn(const std::shared_ptr& engine, - const std::shared_ptr& request) { - std::vector output; - size_t steps = 0; - while (!request->IsTurnComplete()) { - if (++steps >= 100) { - throw std::runtime_error( - "Synthetic windowed request did not complete."); - } - - auto ready = engine->Step(); - if (!ready) { - continue; - } - if (ready != request) { - throw std::runtime_error( - "Synthetic windowed Engine returned an unexpected request."); - } - while (ready->HasUnseenTokens()) { - output.push_back(ready->UnseenToken()); - } - } - return output; -} - -// Wraps the real paged cache only to make reservation lifetime observable. All -// planning, block-table construction, cache binding, and ownership changes -// still run through PagedCacheManager. -class ObservingPagedCacheManager final : public PagedCacheManager { - public: - explicit ObservingPagedCacheManager(std::shared_ptr model) - : PagedCacheManager(std::move(model)) {} - - std::unique_ptr ReserveStep( - const StepPlan& plan) override { - class ObservedReservation final : public CacheStepReservation { - public: - ObservedReservation( - ObservingPagedCacheManager& owner, - std::unique_ptr inner) - : owner_{owner}, inner_{std::move(inner)} {} - - ~ObservedReservation() override { - if (active_) { - active_ = false; - owner_.OnAbandonedReservation(); - } - } - - PagedCacheReservation* PagedReservation() override { - return inner_->PagedReservation(); - } - - void Commit() override { - inner_->Commit(); - if (active_) { - active_ = false; - owner_.OnCommittedReservation(); - } - } - - void Release() override { - inner_->Release(); - if (active_) { - active_ = false; - owner_.OnReleasedReservation(); - } - } - - private: - ObservingPagedCacheManager& owner_; - std::unique_ptr inner_; - bool active_{true}; - }; - - auto reservation = PagedCacheManager::ReserveStep(plan); - ++reservation_count_; - ++active_reservations_; - return std::make_unique( - *this, std::move(reservation)); - } - - size_t ReservationCount() const { return reservation_count_; } - size_t ActiveReservations() const { return active_reservations_; } - size_t CommitCount() const { return commit_count_; } - size_t ReleaseCount() const { return release_count_; } - size_t AbandonedReservationCount() const { - return abandoned_reservation_count_; - } - - private: - void OnCommittedReservation() { - --active_reservations_; - ++commit_count_; - } - - void OnReleasedReservation() { - --active_reservations_; - ++release_count_; - } - - void OnAbandonedReservation() { - --active_reservations_; - ++abandoned_reservation_count_; - } - - size_t reservation_count_{}; - size_t active_reservations_{}; - size_t commit_count_{}; - size_t release_count_{}; - size_t abandoned_reservation_count_{}; -}; - -// Delegates to the production DecoderModelExecutor first. Returning from that -// delegate means the synchronous ORT run completed with the paged cache bound -// as both input and output. Only then does this wrapper observe the real logits -// and raise a retryable failure. -class RetryableAfterRealDecodeExecutor final : public ModelExecutor { - public: - enum class FailurePoint { - None, - BeforeSearchMutation, - AfterSearchMutation, - }; - - explicit RetryableAfterRealDecodeExecutor( - std::unique_ptr inner, - std::shared_ptr cache) - : inner_{std::move(inner)}, cache_{std::move(cache)} {} - - void Decode(ScheduledRequests& scheduled_requests, - ExecutionContext& context) override { - inner_->Decode(scheduled_requests, context); - ++completed_real_decodes_; - - auto failure_point = failure_point_; - if (failure_point == FailurePoint::AfterSearchMutation && - context.plan && !context.plan->requests.empty() && - context.plan->requests.front().target_cache_slots != - failure_target_cache_slots_) { - failure_point = FailurePoint::None; - } else { - failure_point_ = FailurePoint::None; - } - if (failure_point == FailurePoint::None) { - return; - } - - if (!context.plan || context.plan->requests.size() != 1) { - throw std::logic_error( - "Windowed rollback fault expected one planned request."); - } - const auto& entry = context.plan->requests.front(); - failed_unprocessed_token_count_ = entry.unprocessed_token_count; - failed_target_cache_slots_ = entry.target_cache_slots; - if (failure_point == FailurePoint::AfterSearchMutation) { - request_before_failure_ = entry.request->Snapshot(); - cache_before_failure_ = cache_->Snapshot(); - std::vector staged_results; - scheduled_requests.GenerateNextTokensForTransaction( - *context.plan, staged_results); - observed_real_logits_ = - staged_results.size() == scheduled_requests.size(); - staged_sequence_length_ = - entry.request->CurrentSequenceLength(); - staged_search_mutation_ = - staged_sequence_length_ > entry.sequence_length_before; - } else { - const auto logits = scheduled_requests.ProcessLogits(); - observed_real_logits_ = - logits.size() == scheduled_requests.size(); - } - injected_failure_ = true; - throw ModelExecutionError{ - ExecutionFailureKind::RetryableAbort, - "Injected retryable failure after real windowed model execution.", - }; - } - - void FailNextDecode() { - failure_point_ = FailurePoint::BeforeSearchMutation; - } - void FailNextDecodeAfterSearchMutationAtTarget( - size_t target_cache_slots) { - failure_point_ = FailurePoint::AfterSearchMutation; - failure_target_cache_slots_ = target_cache_slots; - } - - bool InjectedFailure() const { return injected_failure_; } - bool ObservedRealLogits() const { return observed_real_logits_; } - size_t CompletedRealDecodes() const { return completed_real_decodes_; } - size_t FailedUnprocessedTokenCount() const { - return failed_unprocessed_token_count_; - } - size_t FailedTargetCacheSlots() const { - return failed_target_cache_slots_; - } - bool StagedSearchMutation() const { return staged_search_mutation_; } - int64_t StagedSequenceLength() const { - return staged_sequence_length_; - } - const RequestStateSnapshot& RequestBeforeFailure() const { - return request_before_failure_; - } - const PagedCacheSnapshot& CacheBeforeFailure() const { - return cache_before_failure_; - } - - private: - std::unique_ptr inner_; - std::shared_ptr cache_; - FailurePoint failure_point_{FailurePoint::None}; - size_t failure_target_cache_slots_{}; - bool injected_failure_{}; - bool observed_real_logits_{}; - bool staged_search_mutation_{}; - size_t completed_real_decodes_{}; - size_t failed_unprocessed_token_count_{}; - size_t failed_target_cache_slots_{}; - int64_t staged_sequence_length_{}; - RequestStateSnapshot request_before_failure_; - PagedCacheSnapshot cache_before_failure_; -}; - -struct FaultInjectingEngine { - std::shared_ptr engine; - std::shared_ptr cache; - RetryableAfterRealDecodeExecutor* executor{}; -}; - -FaultInjectingEngine MakeFaultInjectingEngine( - const std::shared_ptr& model) { - auto cache = - std::make_shared(model); - auto scheduler = Scheduler::Create(model, cache); - auto executor = std::make_unique( - ModelExecutor::Create(model, cache), cache); - auto* executor_observer = executor.get(); - EngineDependencies dependencies{ - cache, std::move(scheduler), std::move(executor)}; - auto engine = - std::make_shared(model, std::move(dependencies)); - return FaultInjectingEngine{ - std::move(engine), std::move(cache), executor_observer}; -} - -void ExpectRequestBlocksEqual(const RequestBlockSnapshot& actual, - const RequestBlockSnapshot& expected) { - EXPECT_EQ(actual.request_id, expected.request_id); - EXPECT_EQ(actual.block_ids, expected.block_ids); - EXPECT_EQ(actual.used_slots, expected.used_slots); - EXPECT_EQ(actual.empty_slots, expected.empty_slots); -} - -void ExpectCacheOwnershipRestored(const PagedCacheSnapshot& actual, - const PagedCacheSnapshot& expected) { - EXPECT_EQ(actual.block_size, expected.block_size); - EXPECT_EQ(actual.total_blocks, expected.total_blocks); - EXPECT_EQ(actual.free_blocks, expected.free_blocks); - EXPECT_EQ(actual.AllocatedBlocks(), expected.AllocatedBlocks()); - EXPECT_TRUE(actual.transaction_reserved_block_ids.empty()); - EXPECT_TRUE(actual.reservations.empty()); - - ASSERT_EQ(actual.requests.size(), expected.requests.size()); - for (size_t i = 0; i < actual.requests.size(); ++i) { - ExpectRequestBlocksEqual(actual.requests[i], expected.requests[i]); - } - - EXPECT_EQ(actual.window_blocks.total_blocks, - expected.window_blocks.total_blocks); - EXPECT_EQ(actual.window_blocks.free_blocks, - expected.window_blocks.free_blocks); - EXPECT_EQ(actual.window_blocks.blocks_per_request, - expected.window_blocks.blocks_per_request); - EXPECT_TRUE( - actual.window_blocks.transaction_reserved_block_ids.empty()); - ASSERT_EQ(actual.window_blocks.requests.size(), - expected.window_blocks.requests.size()); - for (size_t i = 0; i < actual.window_blocks.requests.size(); ++i) { - ExpectRequestBlocksEqual(actual.window_blocks.requests[i], - expected.window_blocks.requests[i]); - } -} - -std::vector RunCleanReplay() { - auto model = LoadWindowedMultiwrapModel(); - auto engine = std::make_shared(model); - std::vector replay_prompt = kInitialPrompt; - replay_prompt.insert(replay_prompt.end(), kFirstTurnOutput.begin(), - kFirstTurnOutput.end()); - replay_prompt.insert(replay_prompt.end(), kContinuation.begin(), - kContinuation.end()); - auto request = MintRequest(*model, replay_prompt); - engine->AddRequest(request); - - auto output = RunTurn(engine, request); - engine->RemoveRequest(request); - return output; -} - -TEST(WindowedTransactionTest, - ContinuedRingWritesRollbackAndRetryMatchCleanReplay) { - const auto clean_replay_output = RunCleanReplay(); - ASSERT_EQ(clean_replay_output, kSecondTurnOutput); - ASSERT_EQ(std::find(clean_replay_output.begin(), - clean_replay_output.end(), - kInvariantFailureToken), - clean_replay_output.end()); - - auto model = LoadWindowedMultiwrapModel(); - auto faulting = MakeFaultInjectingEngine(model); - auto request = MintRequest(*model, kInitialPrompt); - faulting.engine->AddRequest(request); - ASSERT_EQ(RunTurn(faulting.engine, request), kFirstTurnOutput); - ASSERT_EQ(request->Status(), RequestStatus::TurnComplete); - ASSERT_FALSE(request->HasUnseenTokens()); - - request->Continue(kContinuation); - const auto request_before = request->Snapshot(); - const auto cache_before = faulting.cache->Snapshot(); - ASSERT_EQ(request_before.status, RequestStatus::Assigned); - ASSERT_EQ(request_before.current_sequence_length, 12); - ASSERT_EQ(request_before.processed_sequence_length, 3); - ASSERT_EQ(cache_before.requests.size(), 1u); - ASSERT_EQ(cache_before.window_blocks.requests.size(), 1u); - ASSERT_EQ(cache_before.window_blocks.blocks_per_request, 2u); - const size_t ring_period = - cache_before.window_blocks.blocks_per_request * - cache_before.block_size; - ASSERT_EQ(ring_period, 4u); - - const size_t reservations_before = - faulting.cache->ReservationCount(); - const size_t commits_before = faulting.cache->CommitCount(); - const size_t releases_before = faulting.cache->ReleaseCount(); - const size_t decodes_before = - faulting.executor->CompletedRealDecodes(); - faulting.executor->FailNextDecode(); - - try { - static_cast(faulting.engine->Step()); - FAIL() << "Expected retryable failure after real windowed decode."; - } catch (const EngineStepError& error) { - EXPECT_EQ(error.Outcome().kind, - StepOutcomeKind::RetryableBatchAbort); - } - - EXPECT_TRUE(faulting.executor->InjectedFailure()); - EXPECT_TRUE(faulting.executor->ObservedRealLogits()); - EXPECT_EQ(faulting.executor->CompletedRealDecodes(), - decodes_before + 1); - EXPECT_EQ(faulting.executor->FailedUnprocessedTokenCount(), 2u); - EXPECT_EQ(faulting.executor->FailedTargetCacheSlots(), 5u); - // The failed run wrote absolute positions [3, 5), which map to ring - // slots [3, 0]. Thus the injected fault occurred only after a real - // continuation run crossed the four-slot ring boundary. - EXPECT_EQ( - static_cast( - request_before.processed_sequence_length) % - ring_period, - ring_period - 1); - EXPECT_EQ( - (faulting.executor->FailedTargetCacheSlots() - 1) % - ring_period, - 0u); - - const auto request_after = request->Snapshot(); - EXPECT_EQ(request_after.status, RequestStatus::Assigned); - EXPECT_EQ(request_after.current_sequence_length, - request_before.current_sequence_length); - EXPECT_EQ(request_after.processed_sequence_length, - request_before.processed_sequence_length); - EXPECT_FALSE(request->HasUnseenTokens()); - EXPECT_TRUE(faulting.engine->HasPendingRequests()); - - const auto cache_after = faulting.cache->Snapshot(); - ExpectCacheOwnershipRestored(cache_after, cache_before); - EXPECT_EQ(faulting.cache->ReservationCount(), - reservations_before + 1); - EXPECT_EQ(faulting.cache->CommitCount(), commits_before); - EXPECT_EQ(faulting.cache->ReleaseCount(), - releases_before + 1); - EXPECT_EQ(faulting.cache->ActiveReservations(), 0u); - EXPECT_EQ(faulting.cache->AbandonedReservationCount(), 0u); - EXPECT_NO_THROW(ThrowIfInvariantsViolated( - cache_after, std::vector{request_after})); - - const auto retry_output = RunTurn(faulting.engine, request); - EXPECT_EQ(retry_output, clean_replay_output); - EXPECT_EQ(retry_output, kSecondTurnOutput); - EXPECT_EQ(std::find(retry_output.begin(), retry_output.end(), - kInvariantFailureToken), - retry_output.end()); - EXPECT_EQ(request->Status(), RequestStatus::TurnComplete); - EXPECT_EQ(request->CurrentSequenceLength(), 14); - EXPECT_EQ(faulting.cache->ActiveReservations(), 0u); - - faulting.engine->RemoveRequest(request); -} - -TEST(WindowedTransactionTest, - ContinuedRingWritesAndStagedSearchRollbackTogetherBeforeRetry) { - const auto clean_replay_output = RunCleanReplay(); - ASSERT_EQ(clean_replay_output, kSecondTurnOutput); - - auto model = LoadWindowedMultiwrapModel(); - auto faulting = MakeFaultInjectingEngine(model); - auto request = MintRequest(*model, kInitialPrompt); - faulting.engine->AddRequest(request); - ASSERT_EQ(RunTurn(faulting.engine, request), kFirstTurnOutput); - request->Continue(kContinuation); - - ASSERT_EQ(request->Status(), RequestStatus::Assigned); - ASSERT_FALSE(request->HasUnseenTokens()); - - const size_t reservations_before = - faulting.cache->ReservationCount(); - const size_t commits_before = faulting.cache->CommitCount(); - const size_t releases_before = faulting.cache->ReleaseCount(); - // Step commits four two-token continuation chunks before its final one-token prefill. Inject only - // at target 12: that transaction writes absolute position 11 into ring slot 3 after multiple - // wraps, then stages token 28 in Search before raising the retryable failure. - faulting.executor->FailNextDecodeAfterSearchMutationAtTarget(12); - - try { - static_cast(faulting.engine->Step()); - FAIL() << "Expected retryable failure after staged Search mutation."; - } catch (const EngineStepError& error) { - EXPECT_EQ(error.Outcome().kind, - StepOutcomeKind::RetryableBatchAbort); - } - - EXPECT_TRUE(faulting.executor->InjectedFailure()); - EXPECT_TRUE(faulting.executor->ObservedRealLogits()); - EXPECT_TRUE(faulting.executor->StagedSearchMutation()); - EXPECT_EQ(faulting.executor->StagedSequenceLength(), 13); - EXPECT_EQ(faulting.executor->FailedUnprocessedTokenCount(), 1u); - EXPECT_EQ(faulting.executor->FailedTargetCacheSlots(), 12u); - - const auto& request_before = - faulting.executor->RequestBeforeFailure(); - const auto& cache_before = - faulting.executor->CacheBeforeFailure(); - ASSERT_EQ(request_before.status, RequestStatus::Active); - ASSERT_EQ(request_before.current_sequence_length, 12); - ASSERT_EQ(request_before.processed_sequence_length, 11); - - const auto request_after = request->Snapshot(); - EXPECT_EQ(request_after.status, request_before.status); - EXPECT_EQ(request_after.current_sequence_length, - request_before.current_sequence_length); - EXPECT_EQ(request_after.processed_sequence_length, - request_before.processed_sequence_length); - EXPECT_FALSE(request->HasUnseenTokens()); - EXPECT_TRUE(faulting.engine->HasPendingRequests()); - - const auto cache_after = faulting.cache->Snapshot(); - ExpectCacheOwnershipRestored(cache_after, cache_before); - EXPECT_EQ(faulting.cache->ReservationCount(), - reservations_before + 5); - EXPECT_EQ(faulting.cache->CommitCount(), commits_before + 4); - EXPECT_EQ(faulting.cache->ReleaseCount(), - releases_before + 1); - EXPECT_EQ(faulting.cache->ActiveReservations(), 0u); - EXPECT_NO_THROW(ThrowIfInvariantsViolated( - cache_after, std::vector{request_after})); - - const auto retry_output = RunTurn(faulting.engine, request); - EXPECT_EQ(retry_output, clean_replay_output); - EXPECT_EQ(retry_output, kSecondTurnOutput); - EXPECT_EQ(request->Status(), RequestStatus::TurnComplete); - EXPECT_EQ(request->CurrentSequenceLength(), 14); - - faulting.engine->RemoveRequest(request); -} - -} // namespace -} // namespace test -} // namespace Generators diff --git a/test/models/engine/synthetic-windowed-multiwrap/decoder.onnx b/test/models/engine/synthetic-windowed-multiwrap/decoder.onnx deleted file mode 100644 index e15c8e83c476dfc11f51f31328a64883e7f0278d..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 10403 zcma(XYjYIG5xwtPpvUNa0wXRiz$y!L5;7RDV+1nc%0@(rPaS_8`2Ih9xDTV= zHYh{*-5}Ep%n8D_e_5Y6a>G*}eu8yx&>M8!G&uL2#76?5=Q~|LJWbCMm`qYPP8%@g zhi%90c45N343cd!(~1UR+l>Q1af00)l&l+gK5Rj5fItRVAAjj4>1Gv1(`a?-zayit z4XQ|-c7x0+Vme;btyg~u6QS4xPzaiYiK7b@{yWH1(NBrn-%b7a zFh2M=aJ#6FGD>DLNGJH_QBj{~kOPq_O(G+l%zE|6Zx1~Gc+lINg~|ut@3(_qvU-gg zIGt)P+t>O;W#HsS)afKXdTQuU z(~lhwTJRE53!QRn+>}&+a8Wgy~i{@;*4J+v@scR{1b1 z{sIQcG_X%zwCRC^8CO+C=&_#z}#w`ORHH}KrjhIg{LAd3Avhd)uGv5(#I^(Xjg3d(GF%u*ilk*5`WKFU39UN=hU<2$c!Z-0PWgSZXkAPF78_ zX3q5&umV8N2AcKra_OG6lbdCFfUOM$(H6n*_Z5G9#cXmg=C-X9g;$w9FRtC@w* zobLomK%ih6ZEA`V>Mr=IxY%gyxY`L~44Rd7()VLyV9UacH(*7O(I1)1lQc*7+fV>x zR$5TdGZkn+cUUrqRF@jZ)J$Kh&CgJRGSr6@$(l9#Qc0jumB(b}ga3(V$quG>Jsy_c zkzvd(nWy6z?gwcQg$l@1$YDYelFU)$iy%FVFlBEAu1G1ipfG^4G&neB*-<+PPgwxZ zf~p1wR|m*-N?G=3`6ly$ySkctWbWMU4*Y_me~^mqkOdB5Q1u5WV1o&o5Kf|~a03iO zsPc|Pk(y*mKbI)#5LptA!x_n;qW&+e;meRjbP?hzvXgi!a#d<*#xB7k41U0P2WujP zumqE>NeV&x>QhHPE5~43w8DlpTt-A-FJ{Z8S8$T5ap~nUVG$8itw2v@`U3q|HG#A@9w)Sj;JMxzmK15plH3r0d~sz<=zs;u zCMz0<%w|NGyhbuPBLNW|Mrqh`%gDT8;gnf%BcY*0V&!F|KGdqEztz^tB#BXIXSS_%oiXgq5OJE-|3-H`pc^vhQ=w2D%{ekz% z>-x0Tn6DV`QSmJh7#wXA+(fitUN?RM`F?pUFXZ6!1O$1tO2U0B zDV>RMMV2HgbL7u?WvZ6-g^Ka^aTdT_FJB&_SD~ftjMJ3}jLLa$>)a7e7gT zg5Udm17_A(g!em9>^ppyR)rdS8>JAUp)3GkWJuOIXTL3vZcY;&HE7q?PX}(?7Q*E~ z>Z6A-{1q(QCs_9wrE0q7$LUy0W#U^^W3iHf>o_bdf-mt#ItF?XTlKv5-r#F3w>$dAB&7O1#; z9QP{~caK!LIgD4qct<}Rg)wjEiG*$k=loKwE+MdoLkLYf-6)EAh2p-1K^~Gf8iQPQd07P(*=w?~U}EAKjokh>m zbo2Zk^4&1*9ULBmK|_3Ik=apbi*FIKPI)-bWL^1l`7z`MVfsAs+?Es5ml$LoWn{>- zP{y7>NJHJz9zC!UQkN!s$~y^O4nBk;5tz%vk_DAhOrLz<(+A`2`2#S)lF4p#8sKTE zKIj0v`y^i`pF(^+kOrRET+vPNaY-6n#BSed;{>D z8~ch&q*=Kw2`^Pf5QJ6YRehCiRq160E#BNHY*tvme1p@h{LX)V+gxM=tW-7Z{eAOE zWgM9nlJMq+G`O^>w6<(FpE68C%G}(LII@*Evfcd8Xl10bfs9qJZvIQcUz32cLr&wk zu$7#|N%1YGPzqfY{UF+D(Ows=oZlJoJuBKd(awvOh;~7=i=tf;?XqZBM7t{5HPO~Y zyDnPkM7PBE7owFuaB@ef(vn3Wl_@d<*VTtxCwt{@f0kjfGD+rOhFovaFwNnnlaeQSsPQ~fDm;ClVL2=AaP_96GD8#`$ zQ>&#CD*p(|KU%5a?3T2TxTmrBw&9!I4#1 zt^gs_Jlya}fUoODD;uqQpqI21Vw#SvuV@~}>V#PPVzIUrRH0d_lS=anUG~S}+T47M m5BNI4C7!H4xr;CRg3#aMVWHpN8t%_qc>Jb8IKZcXQTTsZTTgoc diff --git a/test/models/engine/synthetic-windowed-multiwrap/genai_config.json b/test/models/engine/synthetic-windowed-multiwrap/genai_config.json deleted file mode 100644 index 79386a0496..0000000000 --- a/test/models/engine/synthetic-windowed-multiwrap/genai_config.json +++ /dev/null @@ -1,56 +0,0 @@ -{ - "model": { - "type": "decoder", - "bos_token_id": 0, - "eos_token_id": 1, - "pad_token_id": 0, - "vocab_size": 32, - "context_length": 16, - "decoder": { - "session_options": { - "log_id": "onnxruntime-genai", - "provider_options": [] - }, - "filename": "decoder.onnx", - "num_attention_heads": 1, - "num_key_value_heads": 1, - "head_size": 1, - "hidden_size": 1, - "num_hidden_layers": 2, - "sliding_window": { - "window_size": 3, - "slide_key_value_cache": false, - "slide_inputs": false, - "layers": [ - 1 - ] - }, - "inputs": { - "input_ids": "input_ids", - "block_table": "block_table", - "block_table_windowed": "block_table_windowed", - "cumulative_sequence_lengths": "cumulative_sequence_lengths", - "past_sequence_lengths": "past_sequence_lengths", - "past_key_names": "past_key_values.%d.key", - "past_value_names": "past_key_values.%d.value" - }, - "outputs": { - "logits": "logits", - "present_key_names": "present.%d.key", - "present_value_names": "present.%d.value" - } - } - }, - "search": { - "max_length": 16, - "chunk_size": 2, - "do_sample": false - }, - "engine": { - "dynamic_batching": { - "block_size": 2, - "num_blocks": 8, - "max_batch_size": 1 - } - } -} diff --git a/test/python/create/create_synthetic_windowed_multiwrap_model.py b/test/python/create/create_synthetic_windowed_multiwrap_model.py deleted file mode 100644 index 8fcf95c134..0000000000 --- a/test/python/create/create_synthetic_windowed_multiwrap_model.py +++ /dev/null @@ -1,631 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -"""Create the deterministic sliding-window Engine continuation fixture. - -The graph models two paged KV-cache layers: - -* layer 0 keeps the full sequence; -* layer 1 uses the runtime's repeated sliding-window block table. - -Both layers store token-and-position encodings. Logits read those encodings back -through both block tables, and an invariant-failure token is selected if the -window table does not repeat, its blocks change across continuation, or its -live cache values disagree with the full cache. -""" - -import argparse -import json -import os - -import numpy as np -import onnx -from onnx import TensorProto, helper, numpy_helper - -VOCAB_SIZE = 32 -BLOCK_SIZE = 2 -WINDOW_SIZE = 3 -CHUNK_SIZE = 2 -MAX_BATCH_SIZE = 1 -RING_BLOCKS = (CHUNK_SIZE + WINDOW_SIZE - 1 + BLOCK_SIZE - 1) // BLOCK_SIZE -NUM_FULL_BLOCKS = 8 -NUM_WINDOW_BLOCKS = RING_BLOCKS * MAX_BATCH_SIZE -CONTEXT_LENGTH = 16 -EOS_TOKEN_ID = 1 -INVARIANT_FAILURE_TOKEN_ID = 31 - - -def _const(name, array): - tensor = numpy_helper.from_array(np.asarray(array)) - tensor.name = name - return tensor - - -def _decoder_graph(): - def i64(value): - return np.asarray(value, dtype=np.int64) - - full_cache_shape = [NUM_FULL_BLOCKS, BLOCK_SIZE, 1, 1] - window_cache_shape = [NUM_WINDOW_BLOCKS, BLOCK_SIZE, 1, 1] - initializers = [ - _const("c0", i64(0)), - _const("c1", i64(1)), - _const("c2", i64(2)), - _const("c3", i64(3)), - _const("c5", i64(5)), - _const("c7", i64(7)), - _const("c13", i64(13)), - _const("c28", i64(28)), - _const("cB", i64(BLOCK_SIZE)), - _const("cR", i64(RING_BLOCKS)), - _const("cEOS", i64(EOS_TOKEN_ID)), - _const("cInvariantFailure", i64(INVARIANT_FAILURE_TOKEN_ID)), - _const("axis0", i64([0])), - _const("axis1", i64([1])), - _const("start1", i64([1])), - _const("end_all", i64([np.iinfo(np.int64).max])), - _const("flat", i64([-1])), - _const("full_cache_shape", i64(full_cache_shape)), - _const("window_cache_shape", i64(window_cache_shape)), - _const( - "vocab_range", - np.arange(VOCAB_SIZE, dtype=np.int64).reshape(1, VOCAB_SIZE), - ), - ] - - nodes = [] - - def node(op_type, inputs, outputs, *, name=None, **attrs): - if name is not None: - attrs["name"] = name - nodes.append(helper.make_node(op_type, inputs, outputs, **attrs)) - - # Resolve each packed token to a request row and an absolute request position. - node("Shape", ["input_ids"], ["ids_shape"]) - node("Squeeze", ["ids_shape"], ["num_tokens"]) - node("Range", ["c0", "num_tokens", "c1"], ["token_index"]) - node( - "Slice", - ["cumulative_sequence_lengths", "start1", "end_all", "axis0"], - ["boundaries_i32"], - ) - node("Cast", ["boundaries_i32"], ["boundaries"], to=TensorProto.INT64) - node("Unsqueeze", ["token_index", "axis1"], ["token_index_col"]) - node("Unsqueeze", ["boundaries", "axis0"], ["boundaries_row"]) - node("GreaterOrEqual", ["token_index_col", "boundaries_row"], ["at_or_past"]) - node("Cast", ["at_or_past"], ["at_or_past_i64"], to=TensorProto.INT64) - node("ReduceSum", ["at_or_past_i64", "axis1"], ["row_id"], keepdims=0) - - node( - "Cast", - ["cumulative_sequence_lengths"], - ["cumulative_sequence_lengths_i64"], - to=TensorProto.INT64, - ) - node( - "Gather", - ["cumulative_sequence_lengths_i64", "row_id"], - ["row_start"], - axis=0, - ) - node("Sub", ["token_index", "row_start"], ["offset_in_row"]) - node( - "Cast", - ["past_sequence_lengths"], - ["past_sequence_lengths_i64"], - to=TensorProto.INT64, - ) - node( - "Gather", - ["past_sequence_lengths_i64", "row_id"], - ["past_of_row"], - axis=0, - ) - node("Add", ["past_of_row", "offset_in_row"], ["pos"]) - node("Sub", ["pos", "c1"], ["previous_pos_unclamped"]) - node("Max", ["previous_pos_unclamped", "c0"], ["previous_pos"]) - - node("Cast", ["block_table"], ["block_table_i64"], to=TensorProto.INT64) - node( - "Cast", - ["block_table_windowed"], - ["block_table_windowed_i64"], - to=TensorProto.INT64, - ) - node("Unsqueeze", ["row_id", "axis1"], ["row_id_col"]) - - def map_positions(prefix, positions, table): - block_col = f"{prefix}_block_col" - block_col_base = f"{prefix}_block_col_base" - slot = f"{prefix}_slot_in_block" - block_col_col = f"{prefix}_block_col_col" - gather_index = f"{prefix}_block_gather_index" - block_id = f"{prefix}_block_id" - block_base = f"{prefix}_block_base" - physical = f"{prefix}_physical" - - node("Div", [positions, "cB"], [block_col]) - node("Mul", [block_col, "cB"], [block_col_base]) - node("Sub", [positions, block_col_base], [slot]) - node("Unsqueeze", [block_col, "axis1"], [block_col_col]) - node("Concat", ["row_id_col", block_col_col], [gather_index], axis=1) - node("GatherND", [table, gather_index], [block_id]) - node("Mul", [block_id, "cB"], [block_base]) - node("Add", [block_base, slot], [physical]) - return block_col, block_id, physical - - current_block_col, _, current_full_physical = map_positions("current_full", "pos", "block_table_i64") - _, current_window_block, current_window_physical = map_positions( - "current_window", "pos", "block_table_windowed_i64" - ) - _, _, previous_full_physical = map_positions("previous_full", "previous_pos", "block_table_i64") - _, _, previous_window_physical = map_positions("previous_window", "previous_pos", "block_table_windowed_i64") - - # Position zero remains in the full cache. Its value records the first - # window block id so continuation can prove that the ring stayed resident. - node("Gather", ["block_table_i64", "c0"], ["first_full_block_per_row"], axis=1) - node( - "Gather", - ["first_full_block_per_row", "row_id"], - ["first_full_block"], - axis=0, - ) - node("Mul", ["first_full_block", "cB"], ["first_full_physical"]) - node( - "Gather", - ["block_table_windowed_i64", "c0"], - ["first_window_block_per_row"], - axis=1, - ) - node( - "Gather", - ["first_window_block_per_row", "row_id"], - ["first_window_block"], - axis=0, - ) - - # Store exact integer-valued float encodings. Including the absolute - # position makes a continuation that resets past_sequence_lengths diverge. - node("Mul", ["input_ids", "c7"], ["key_token_term"]) - node("Mul", ["pos", "c3"], ["key_position_term"]) - node("Add", ["key_token_term", "key_position_term"], ["key_without_bias"]) - node("Add", ["key_without_bias", "c1"], ["key_encoding_i64"]) - node("Cast", ["key_encoding_i64"], ["key_encoding"], to=TensorProto.FLOAT) - - node("Mul", ["input_ids", "c5"], ["value_token_term"]) - node("Mul", ["pos", "c2"], ["value_position_term"]) - node( - "Add", - ["value_token_term", "value_position_term"], - ["value_without_bias"], - ) - node("Add", ["value_without_bias", "c2"], ["value_encoding_i64"]) - node( - "Cast", - ["value_encoding_i64"], - ["window_value_encoding"], - to=TensorProto.FLOAT, - ) - node( - "Cast", - [current_window_block], - ["window_owner_encoding"], - to=TensorProto.FLOAT, - ) - - node( - "Reshape", - ["past_key_values.0.key", "flat"], - ["past_full_key_flat"], - ) - node( - "Reshape", - ["past_key_values.0.value", "flat"], - ["past_full_value_flat"], - ) - node( - "Reshape", - ["past_key_values.1.key", "flat"], - ["past_window_key_flat"], - ) - node( - "Reshape", - ["past_key_values.1.value", "flat"], - ["past_window_value_flat"], - ) - node( - "Unsqueeze", - [current_full_physical, "axis1"], - ["current_full_scatter_index"], - ) - node( - "Unsqueeze", - [current_window_physical, "axis1"], - ["current_window_scatter_index"], - ) - - node( - "ScatterND", - ["past_full_key_flat", "current_full_scatter_index", "key_encoding"], - ["present_full_key_flat"], - ) - node( - "ScatterND", - [ - "past_full_value_flat", - "current_full_scatter_index", - "window_owner_encoding", - ], - ["present_full_value_flat"], - ) - node( - "ScatterND", - ["past_window_key_flat", "current_window_scatter_index", "key_encoding"], - ["present_window_key_flat"], - ) - node( - "ScatterND", - [ - "past_window_value_flat", - "current_window_scatter_index", - "window_value_encoding", - ], - ["present_window_value_flat"], - ) - - node( - "Reshape", - ["present_full_key_flat", "full_cache_shape"], - ["present.0.key"], - ) - node( - "Reshape", - ["present_full_value_flat", "full_cache_shape"], - ["present.0.value"], - ) - node( - "Reshape", - ["present_window_key_flat", "window_cache_shape"], - ["present.1.key"], - ) - node( - "Reshape", - ["present_window_value_flat", "window_cache_shape"], - ["present.1.value"], - ) - - # Logits consume values read through both cache layers, including the - # previous live window position after the ring has wrapped. - node( - "Gather", - ["present_full_key_flat", "first_full_physical"], - ["read_full_first_key"], - axis=0, - name="read_full_first_key", - ) - node( - "Gather", - ["present_full_key_flat", current_full_physical], - ["read_full_current_key"], - axis=0, - name="read_full_current_key", - ) - node( - "Gather", - ["present_full_key_flat", previous_full_physical], - ["read_full_previous_key"], - axis=0, - name="read_full_previous_key", - ) - node( - "Gather", - ["present_full_value_flat", "first_full_physical"], - ["read_first_window_owner"], - axis=0, - name="read_first_window_owner", - ) - node( - "Gather", - ["present_window_key_flat", previous_window_physical], - ["read_window_previous_key"], - axis=0, - name="read_window_previous_key", - ) - node( - "Gather", - ["present_window_key_flat", current_window_physical], - ["read_window_current_key"], - axis=0, - name="read_window_current_key", - ) - node( - "Gather", - ["present_window_value_flat", previous_window_physical], - ["read_window_previous_value"], - axis=0, - name="read_window_previous_value", - ) - node( - "Gather", - ["present_window_value_flat", current_window_physical], - ["read_window_current_value"], - axis=0, - name="read_window_current_value", - ) - - # The current and prior-cycle columns must name the same physical block. - node("Sub", [current_block_col, "cR"], ["prior_cycle_col_unclamped"]) - node("Max", ["prior_cycle_col_unclamped", "c0"], ["prior_cycle_col"]) - node("Unsqueeze", ["prior_cycle_col", "axis1"], ["prior_cycle_col_col"]) - node( - "Concat", - ["row_id_col", "prior_cycle_col_col"], - ["prior_cycle_gather_index"], - axis=1, - ) - node( - "GatherND", - ["block_table_windowed_i64", "prior_cycle_gather_index"], - ["prior_cycle_window_block"], - ) - node( - "GreaterOrEqual", - [current_block_col, "cR"], - ["has_prior_block_cycle"], - ) - node( - "Equal", - [current_window_block, "prior_cycle_window_block"], - ["window_block_repeats"], - ) - node("Not", ["has_prior_block_cycle"], ["before_first_block_cycle"]) - node( - "Or", - ["before_first_block_cycle", "window_block_repeats"], - ["repeated_window_block_valid"], - name="guard_repeated_window_block", - ) - - node( - "Equal", - ["read_full_previous_key", "read_window_previous_key"], - ["previous_cache_values_match"], - ) - node( - "Equal", - ["read_full_current_key", "read_window_current_key"], - ["current_cache_values_match"], - ) - node( - "Cast", - ["first_window_block"], - ["first_window_block_f"], - to=TensorProto.FLOAT, - ) - node( - "Equal", - ["read_first_window_owner", "first_window_block_f"], - ["window_owner_stable"], - name="guard_stable_window_owner", - ) - node( - "And", - ["previous_cache_values_match", "current_cache_values_match"], - ["cache_values_match"], - ) - node( - "And", - ["cache_values_match", "window_owner_stable"], - ["cache_and_owner_valid"], - ) - node( - "And", - ["cache_and_owner_valid", "repeated_window_block_valid"], - ["window_invariants_valid"], - ) - - score_terms = [ - "read_full_first_key", - "read_full_current_key", - "read_window_previous_key", - "read_window_current_key", - "read_window_previous_value", - "read_window_current_value", - ] - score = score_terms[0] - for index, term in enumerate(score_terms[1:], start=1): - output = f"score_sum_{index}" - node("Add", [score, term], [output]) - score = output - node("Cast", [score], ["score_i64"], to=TensorProto.INT64) - node("Div", ["score_i64", "c28"], ["score_div"]) - node("Mul", ["score_div", "c28"], ["score_floor"]) - node("Sub", ["score_i64", "score_floor"], ["score_mod"]) - node("Add", ["score_mod", "c2"], ["normal_next_token"]) - node( - "Where", - ["window_invariants_valid", "normal_next_token", "cInvariantFailure"], - ["guarded_next_token"], - ) - - # EOS at absolute positions 2 and 13 creates two short turns while leaving - # max_length headroom. The continuation spans positions 3..11; normal - # outputs at 11 and 12 validate both columns of the two-block ring. - node("Equal", ["pos", "c2"], ["is_first_turn_eos_position"]) - node("Equal", ["pos", "c13"], ["is_second_turn_eos_position"]) - node( - "Or", - ["is_first_turn_eos_position", "is_second_turn_eos_position"], - ["is_eos_position"], - ) - node( - "Where", - ["is_eos_position", "cEOS", "guarded_next_token"], - ["next_token"], - ) - - node("Unsqueeze", ["next_token", "axis1"], ["next_token_col"]) - node("Equal", ["next_token_col", "vocab_range"], ["is_next_per_token"]) - node("Sub", ["boundaries", "c1"], ["last_token_index"]) - node( - "Gather", - ["is_next_per_token", "last_token_index"], - ["is_next_per_request"], - axis=0, - ) - node( - "Cast", - ["is_next_per_request"], - ["logits"], - to=TensorProto.FLOAT16, - ) - - inputs = [ - helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["num_tokens"]), - helper.make_tensor_value_info( - "cumulative_sequence_lengths", - TensorProto.INT32, - ["batch_plus_1"], - ), - helper.make_tensor_value_info("past_sequence_lengths", TensorProto.INT32, ["batch"]), - helper.make_tensor_value_info("block_table", TensorProto.INT32, ["batch", "max_blocks"]), - helper.make_tensor_value_info( - "block_table_windowed", - TensorProto.INT32, - ["batch", "max_blocks"], - ), - helper.make_tensor_value_info( - "past_key_values.0.key", - TensorProto.FLOAT, - full_cache_shape, - ), - helper.make_tensor_value_info( - "past_key_values.0.value", - TensorProto.FLOAT, - full_cache_shape, - ), - helper.make_tensor_value_info( - "past_key_values.1.key", - TensorProto.FLOAT, - window_cache_shape, - ), - helper.make_tensor_value_info( - "past_key_values.1.value", - TensorProto.FLOAT, - window_cache_shape, - ), - ] - outputs = [ - helper.make_tensor_value_info("logits", TensorProto.FLOAT16, ["batch_size", VOCAB_SIZE]), - helper.make_tensor_value_info("present.0.key", TensorProto.FLOAT, full_cache_shape), - helper.make_tensor_value_info("present.0.value", TensorProto.FLOAT, full_cache_shape), - helper.make_tensor_value_info("present.1.key", TensorProto.FLOAT, window_cache_shape), - helper.make_tensor_value_info("present.1.value", TensorProto.FLOAT, window_cache_shape), - ] - return helper.make_graph( - nodes, - "synthetic_windowed_multiwrap_decoder", - inputs, - outputs, - initializer=initializers, - ) - - -def create_decoder(output_dir): - model = helper.make_model( - _decoder_graph(), - opset_imports=[helper.make_operatorsetid("", 17)], - ir_version=9, - producer_name="onnxruntime-genai", - producer_version="0.0.0", - ) - metadata = model.metadata_props.add() - metadata.key = "fixture" - metadata.value = "engine-windowed-multiwrap-continuation" - onnx.checker.check_model(model) - onnx.save_model(model, os.path.join(output_dir, "decoder.onnx")) - - -def create_config(output_dir): - config = { - "model": { - "type": "decoder", - "bos_token_id": 0, - "eos_token_id": EOS_TOKEN_ID, - "pad_token_id": 0, - "vocab_size": VOCAB_SIZE, - "context_length": CONTEXT_LENGTH, - "decoder": { - "session_options": { - "log_id": "onnxruntime-genai", - "provider_options": [], - }, - "filename": "decoder.onnx", - "num_attention_heads": 1, - "num_key_value_heads": 1, - "head_size": 1, - "hidden_size": 1, - "num_hidden_layers": 2, - "sliding_window": { - "window_size": WINDOW_SIZE, - "slide_key_value_cache": False, - "slide_inputs": False, - "layers": [1], - }, - "inputs": { - "input_ids": "input_ids", - "block_table": "block_table", - "block_table_windowed": "block_table_windowed", - "cumulative_sequence_lengths": "cumulative_sequence_lengths", - "past_sequence_lengths": "past_sequence_lengths", - "past_key_names": "past_key_values.%d.key", - "past_value_names": "past_key_values.%d.value", - }, - "outputs": { - "logits": "logits", - "present_key_names": "present.%d.key", - "present_value_names": "present.%d.value", - }, - }, - }, - "search": { - "max_length": CONTEXT_LENGTH, - "chunk_size": CHUNK_SIZE, - "do_sample": False, - }, - "engine": { - "dynamic_batching": { - "block_size": BLOCK_SIZE, - "num_blocks": NUM_FULL_BLOCKS, - "max_batch_size": MAX_BATCH_SIZE, - }, - }, - } - with open(os.path.join(output_dir, "genai_config.json"), "w") as config_file: - json.dump(config, config_file, indent=2) - config_file.write("\n") - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--output_dir", - default=os.path.join( - os.path.dirname(__file__), - "..", - "..", - "models", - "engine", - "synthetic-windowed-multiwrap", - ), - ) - args = parser.parse_args() - output_dir = os.path.normpath(args.output_dir) - os.makedirs(output_dir, exist_ok=True) - create_decoder(output_dir) - create_config(output_dir) - - -if __name__ == "__main__": - main() diff --git a/test/python/models/test_engine_windowed_multiwrap.py b/test/python/models/test_engine_windowed_multiwrap.py deleted file mode 100644 index dda907a08f..0000000000 --- a/test/python/models/test_engine_windowed_multiwrap.py +++ /dev/null @@ -1,199 +0,0 @@ -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. - -"""Executable Engine coverage for continuation across window-ring wraps.""" - -import json -from itertools import pairwise -from pathlib import Path - -import numpy as np -import onnx -import onnxruntime_genai as og -import pytest - -_MODEL_SUBPATH = Path("engine") / "synthetic-windowed-multiwrap" - -_VOCAB_SIZE = 32 -_BLOCK_SIZE = 2 -_WINDOW_SIZE = 3 -_CHUNK_SIZE = 2 -_RING_BLOCKS = 2 -_RING_PERIOD = _RING_BLOCKS * _BLOCK_SIZE -_NUM_FULL_BLOCKS = 8 -_MAX_LENGTH = 16 -_EOS_TOKEN_ID = 1 -_INVARIANT_FAILURE_TOKEN_ID = 31 - -_INITIAL_PROMPT = [4, 6] -_EXPECTED_FIRST_TURN = [12] -_CONTINUATION = [8, 3, 10, 5, 12, 7, 14, 9, 16] -_EXPECTED_SECOND_TURN = [28, 17] - -_DEVICES = ["cpu"] + (["cuda"] if og.is_cuda_available() else []) - - -def _fixture_path(test_data_path) -> Path: - if not test_data_path: - pytest.skip("--test_models is required for the synthetic Engine fixture") - path = Path(test_data_path) / _MODEL_SUBPATH - if not path.exists(): - pytest.fail(f"synthetic Engine fixture is missing: {path}") - return path - - -def _cache_key(token, position): - return token * 7 + position * 3 + 1 - - -def _cache_value(token, position): - return token * 5 + position * 2 + 2 - - -def _normal_token(sequence, position): - """Mirror the graph's cache-read score for a non-EOS position.""" - first_key = _cache_key(sequence[0], 0) - previous_key = _cache_key(sequence[position - 1], position - 1) - current_key = _cache_key(sequence[position], position) - previous_value = _cache_value(sequence[position - 1], position - 1) - current_value = _cache_value(sequence[position], position) - score = first_key + current_key + previous_key + current_key + previous_value + current_value - return score % 28 + 2 - - -def _wrap_count(start_position, token_count): - positions = range(start_position, start_position + token_count) - slots = [position % _RING_PERIOD for position in positions] - return sum(current < previous for previous, current in pairwise(slots)) - - -def _new_request(engine, model, prompt): - params = og.GeneratorParams(model) - params.set_search_options(do_sample=False, max_length=_MAX_LENGTH) - request = og.Request(params) - request.add_tokens(np.asarray(prompt, dtype=np.int32)) - engine.add_request(request) - return request - - -def _run_turn(engine, request): - tokens = [] - steps = 0 - while not request.is_turn_complete(): - ready = engine.step() - steps += 1 - assert steps < 100, "synthetic turn did not reach its absolute-position EOS" - if ready is None: - # Partial prefill chunks commit cache progress without sampling. - continue - assert ready is request - while ready.has_unseen_tokens(): - tokens.append(ready.get_unseen_token()) - return tokens - - -def test_windowed_multiwrap_fixture_schema(test_data_path): - model_path = _fixture_path(test_data_path) - config = json.loads((model_path / "genai_config.json").read_text(encoding="utf-8")) - decoder = config["model"]["decoder"] - dynamic = config["engine"]["dynamic_batching"] - - assert decoder["num_hidden_layers"] == 2 - assert decoder["sliding_window"] == { - "window_size": _WINDOW_SIZE, - "slide_key_value_cache": False, - "slide_inputs": False, - "layers": [1], - } - assert decoder["inputs"]["block_table_windowed"] == "block_table_windowed" - assert config["model"]["vocab_size"] == _VOCAB_SIZE - assert config["model"]["eos_token_id"] == _EOS_TOKEN_ID - assert config["search"]["max_length"] == _MAX_LENGTH - assert config["search"]["chunk_size"] == _CHUNK_SIZE - assert dynamic["block_size"] == _BLOCK_SIZE - assert dynamic["num_blocks"] == _NUM_FULL_BLOCKS - assert dynamic["max_batch_size"] == 1 - - assert (_CHUNK_SIZE + _WINDOW_SIZE - 1 + _BLOCK_SIZE - 1) // _BLOCK_SIZE == _RING_BLOCKS - continuation_start = len(_INITIAL_PROMPT) + len(_EXPECTED_FIRST_TURN) - assert len(_CONTINUATION) > 2 * _RING_PERIOD - assert _wrap_count(continuation_start, len(_CONTINUATION)) == 2 - - graph_model = onnx.load(model_path / "decoder.onnx", load_external_data=False) - onnx.checker.check_model(graph_model) - metadata = {entry.key: entry.value for entry in graph_model.metadata_props} - assert metadata["fixture"] == "engine-windowed-multiwrap-continuation" - - input_shapes = { - value.name: [dimension.dim_value for dimension in value.type.tensor_type.shape.dim] - for value in graph_model.graph.input - if value.name.startswith("past_key_values") - } - assert input_shapes["past_key_values.0.key"] == [_NUM_FULL_BLOCKS, 2, 1, 1] - assert input_shapes["past_key_values.0.value"] == [_NUM_FULL_BLOCKS, 2, 1, 1] - assert input_shapes["past_key_values.1.key"] == [2, 2, 1, 1] - assert input_shapes["past_key_values.1.value"] == [2, 2, 1, 1] - - node_names = {node.name for node in graph_model.graph.node} - assert { - "read_full_current_key", - "read_full_previous_key", - "read_window_previous_key", - "read_window_current_value", - "guard_repeated_window_block", - "guard_stable_window_owner", - } <= node_names - - -@pytest.mark.parametrize("device", _DEVICES) -def test_continuation_crosses_two_window_ring_wraps_and_matches_clean_replay(test_data_path, device): - model_path = _fixture_path(test_data_path) - config = og.Config(str(model_path)) - config.clear_providers() - if device == "cuda": - config.append_provider("cuda") - model = og.Model(config) - - # The first normal result is followed by EOS at absolute position 2, well - # below the session max. EOS is not part of the retained logical sequence. - assert _normal_token(_INITIAL_PROMPT, 1) == _EXPECTED_FIRST_TURN[0] - engine = og.Engine(model) - request = _new_request(engine, model, _INITIAL_PROMPT) - first_turn = _run_turn(engine, request) - - assert first_turn == _EXPECTED_FIRST_TURN - assert request.is_turn_complete() - first_turn_length = len(_INITIAL_PROMPT) + len(first_turn) - assert first_turn_length == 3 - assert first_turn_length < _MAX_LENGTH - - # Positions 3..11 traverse the four-slot ring twice. The graph returns 31 - # if the repeated table, retained ring ownership, or values read through - # the full/windowed caches disagree. Tokens at positions 11 and 12 check - # both columns of the two-block ring before EOS at position 13. - request.continue_with(np.asarray(_CONTINUATION, dtype=np.int32)) - assert not request.is_turn_complete() - second_turn = _run_turn(engine, request) - - assert second_turn == _EXPECTED_SECOND_TURN - assert _INVARIANT_FAILURE_TOKEN_ID not in second_turn - assert request.is_turn_complete() - retained_length = first_turn_length + len(_CONTINUATION) + len(second_turn) - assert retained_length == 14 - assert retained_length < _MAX_LENGTH - engine.remove_request(request) - - # Recompute the same logical context in a fresh cache. Exact parity proves - # that continuation's retained full cache and twice-wrapped window cache - # agree with a clean chunked replay at absolute positions 0..13. - replay_prompt = _INITIAL_PROMPT + first_turn + _CONTINUATION - assert _normal_token(replay_prompt, len(replay_prompt) - 1) == _EXPECTED_SECOND_TURN[0] - replay_after_first_output = replay_prompt + _EXPECTED_SECOND_TURN[:1] - assert _normal_token(replay_after_first_output, len(replay_after_first_output) - 1) == _EXPECTED_SECOND_TURN[1] - replay_engine = og.Engine(model) - replay_request = _new_request(replay_engine, model, replay_prompt) - replay_turn = _run_turn(replay_engine, replay_request) - - assert replay_turn == second_turn == _EXPECTED_SECOND_TURN - assert _INVARIANT_FAILURE_TOKEN_ID not in replay_turn - replay_engine.remove_request(replay_request)