Harden Engine request lifetime and sampler acquisition - #2448
Draft
bmehta001 wants to merge 23 commits into
Draft
Conversation
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
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
Preserve the upstream per-request logits optimization while retaining scheduled-token accounting for continuous decoding in VarlenDecoderIO. Conflict resolved: src/engine/decoders/varlen_decoder_io.cpp Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: b232f82c-f25c-429a-a855-7f8c8f40bbf3
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
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
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
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
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
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
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>
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>
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>
…ing' into bmehta001-engine-admission-invariants-hardening
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>
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
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
Integrate the latest main history while retaining only the validated admission, invariant, rollback, lifetime, and sampler hardening follow-ups. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e54c6195-e3c4-4467-abcf-72bdff3260dc
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
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
Integrate current main benchmark, guidance, model-state, and paged-cache changes while preserving only #2448 lifetime, admission, diagnostics, sampler, and windowed rollback hardening. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e54c6195-e3c4-4467-abcf-72bdff3260dc
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
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
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>
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Request/Engine/scheduler friendships using narrow internal ownership and continuation capabilities; no public scheduler status or mutating abandonment callbacks are addedsrc/cuda/interface.cppScope note
The final
main...HEADdiff contains no DML replay implementation, test, documentation claim, or public API. An old DML commit remains only in inherited branch history from PR #2423 and has no effect on this PR's review diff.The CUDA sampler index helper is intentionally implementation-local; this PR does not add a standalone pool header or dedicated pool test target.
Validation
clang-format 20.1.0 --dry-run -Werror: passed across all 386 tracked C/C++ filesEngineUnitTests: passedUnitTests: passedUSE_CXX17syntax check: passed/std:c++20 /Qspectre /W4 /WX /Zssyntax check of changed Engine sources: passedNo local CUDA compiler is installed, so the provider-specific CUDA translation-unit build is delegated to CI.
CI notes
Earlier
lint-cppand Windows CUDA failures were self-hosted runner disconnects, not format or build failures. The two earlier failed ADO integration jobs exited with code 127 before tests began. Other prior Windows CPU, DirectML, WinML, WebGPU, Linux CPU, Linux ARM64, and Linux CUDA jobs passed. The latest push retriggers CI.A full local Windows MSBuild remains blocked before project compilation by missing MSVC Spectre-mitigated libraries (
MSB8040). Security/Spectre settings were not weakened.