Enable model-selected Engine chat batching - #1071
Conversation
Use each model's genai_config batching contract so independent chat sessions can share an ORT GenAI Engine while Generator models retain their existing behavior. Files changed: - Add the owner-thread Engine dispatcher and per-session adapter. - Route chat creation, continuation, cancellation, usage, and option changes through the selected backend. - Parse and validate static and dynamic Engine configuration. - Add configuration, template, concurrency, and search-option coverage. - Pin ORT GenAI and packaging to 0.15.3. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9b90e66e-af55-404d-b606-7899f6f73c54
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Remove compatibility branches for OGA releases without Engine support so missing APIs fail during compilation instead of at model load. Files changed: - sdk_v2/cpp/CMakeLists.txt - sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc - sdk_v2/cpp/src/inferencing/generative/genai_model_instance.cc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9b90e66e-af55-404d-b606-7899f6f73c54
Preserve the OGA 0.15.3 release requirement while adopting main's nightly GenAI validation path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9b90e66e-af55-404d-b606-7899f6f73c54
Retain OGA 0.15.2 until the next Engine-capable stable package is published, while non-release CI continues using the selected nightly. Remove optional OGA hooks from the shared test entry point because not every test target consumes OGA. Files changed: - .pipelines/foundry-local-packaging.yml - .pipelines/v2/sdk_v2-pipeline-plan.md - sdk_v2/deps_versions.json - sdk_v2/cpp/test/test_main.cc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9b90e66e-af55-404d-b606-7899f6f73c54
The shared test entry point is compiled by targets with and without OGA include paths. Retain header detection so OGA-linked tests shut down cleanly without imposing that dependency on cache-only tests. Files changed: - sdk_v2/cpp/test/test_main.cc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9b90e66e-af55-404d-b606-7899f6f73c54
There was a problem hiding this comment.
🟡 Changes recommended
Release dependencies cannot compile the new API, and Engine failure and finish-reason handling contain correctness issues.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds model-configured static or dynamic ORT GenAI Engine batching while preserving the existing Generator backend.
Changes:
- Parses Engine batching configuration and selects the appropriate backend.
- Adds a model-owned threaded Engine dispatcher with retained dynamic conversations.
- Extends chat generation, cancellation, usage accounting, templates, and tests.
File summaries
| File | Description |
|---|---|
sdk_v2/cpp/CMakeLists.txt |
Registers Engine sources. |
sdk_v2/cpp/src/inferencing/generative/genai_config.h |
Defines Engine configuration. |
sdk_v2/cpp/src/inferencing/generative/genai_config.cc |
Parses and validates batching modes. |
sdk_v2/cpp/src/inferencing/generative/genai_model_instance.h |
Stores the model-owned Engine. |
sdk_v2/cpp/src/inferencing/generative/genai_model_instance.cc |
Creates and destroys the Engine. |
sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h |
Generalizes retained generator operations. |
sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.cc |
Adds optional turn usage. |
sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_engine.h |
Declares the Engine dispatcher. |
sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_engine.cc |
Implements scheduling and event routing. |
sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_generator.h |
Declares the Engine adapter. |
sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_generator.cc |
Implements Engine-backed generation. |
sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.h |
Adapts the Generator interface. |
sdk_v2/cpp/src/inferencing/generative/chat/onnx_chat_generator.cc |
Shares guidance handling. |
sdk_v2/cpp/src/inferencing/generative/chat/search_options.h |
Adds retained-setting and limit helpers. |
sdk_v2/cpp/src/inferencing/generative/chat/search_options.cc |
Resolves limits and guidance. |
sdk_v2/cpp/src/inferencing/generative/chat/chat_template.h |
Declares continuation prompt support. |
sdk_v2/cpp/src/inferencing/generative/chat/chat_template.cc |
Derives assistant-turn boundaries. |
sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h |
Stores backend-neutral retained state. |
sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc |
Selects and manages chat backends. |
sdk_v2/cpp/test/test_main.cc |
Shuts down ORT GenAI after tests. |
sdk_v2/cpp/test/internal_api/genai_config_test.cc |
Tests Engine configuration parsing. |
sdk_v2/cpp/test/internal_api/chat/search_options_test.cc |
Tests limits and retained settings. |
sdk_v2/cpp/test/internal_api/chat/chat_template_test.cc |
Tests continuation boundaries. |
sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc |
Adds concurrent-session coverage. |
Review details
- Files reviewed: 24/24 changed files
- Comments generated: 4
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| std::optional<ChatTurnUsage> OnnxEngineChatGenerator::GetTurnUsage() const { | ||
| const auto result = engine_.GetTurnResult(conversation_); | ||
| return ChatTurnUsage{ | ||
| static_cast<int>(result.prompt_tokens + result.cached_prompt_tokens), | ||
| static_cast<int>(result.generated_tokens), | ||
| }; |
| FL_LOG_AND_THROW(logger, FOUNDRY_LOCAL_ERROR_INTERNAL, | ||
| "model ", model_id_, " declares an Engine backend, but Engine is not supported for multimodal models"); |
Merge origin/main to resolve the PR conflict without rewriting history. Engine events: fail all conversations when an engine-level failure has no request handle; propagate native finish reasons so implicit output and session limits report LENGTH. C++ cleanup: wrap the long Engine validation message and remove duplicate gtest includes. Files changed: merged main updates; Engine chat dispatcher/generator/session metadata; GenAI model instance; test main. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
The pinned GenAI dependency cannot compile the new API, and request limiting, undo, and lifetime paths have correctness issues.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 25/25 changed files
- Comments generated: 5
- Review effort level: Balanced
| OnnxEngineChatGenerator::~OnnxEngineChatGenerator() { | ||
| try { | ||
| engine_.Close(conversation_); | ||
| } catch (...) { |
| return OnnxEngineChatGenerator::Create(messages, options, model, tool_ctx); | ||
| } | ||
|
|
||
| return OnnxChatGenerator::Create(messages, options, model, tool_ctx, /*use_full_context=*/true); |
| if (cached_generator_ && | ||
| !cached_search_options_.HasSameRetainedGenerationSettings(effective_options)) { | ||
| cached_generator_.reset(); | ||
| cached_tool_ctx_ = {}; | ||
| } |
| EXPECT_EQ(session.GetHistory()[1].GetSimpleText(), text); | ||
| } | ||
|
|
||
| TEST_F(ChatSessionTest, ConcurrentIndependentSessions) { |
Destroy retained generators before releasing model ownership, preserve undo correctness after generator rebuilds, and keep the Generator path's implicit limit behavior unchanged. Use bounded Generator context for one-shot OpenAI JSON requests and map the CI nightly's native StopSequence finish reason. Stage the shared chat model with dynamic Engine batching so concurrency and cancellation coverage exercise the dispatcher. Files changed: chat_session.cc, chat_session.h, onnx_engine_chat_generator.cc, chat_session_test.cc. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
The current dependency pin and finish-reason symbol prevent a clean build, with additional retained-state correctness and cleanup issues.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc:515
- Retained requests are rebuilt only when the boolean guidance mode changes, but
ApplyGuidanceOptionsalso bakes the specificguidance_typeandguidance_datainto the request. Two tool-only turns with different response schemas therefore reuse the first turn's grammar and constrain the second response incorrectly. Include changes to those guidance values in the invalidation condition.
sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_generator.cc:172 - After
CreateConversationinserts the native request into the engine map, any exception fromBeginTurn, tokenizer-stream creation, or generator allocation exits withoutClose. The map retains the conversation and request indefinitely, potentially consuming batching/KV capacity after a failed request. Add failure cleanup until ownership has transferred to the returned generator.
- Files reviewed: 25/25 changed files
- Comments generated: 2
- Review effort level: Balanced
| engine_ = OgaEngine::Create(model_.GetOgaModel()); | ||
| event_buffer_ = engine_->CreateEventBuffer(model_.GetGenAIConfig().EngineMaxBatchSize().value_or(1) * 2); |
There was a problem hiding this comment.
The release pin intentionally remains 0.15.2 because 0.15.3 is not published: the NuGet gallery currently lists 0.15.2 as latest, and the configured AIFoundryLocal public feed returns HTTP 404 for 0.15.3. Non-release CI builds this PR against the selected Engine-capable 0.16.0-dev1001400138 nightly (.pipelines/foundry-local-packaging.yml:64-66), while the stable pin will move when an Engine-capable release is available. No dependency-pin change is possible in this PR without referencing a nonexistent package.
| std::optional<flFinishReason> MapFinishReason(OgaFinishReason reason) { | ||
| switch (reason) { | ||
| case OgaFinishReason_Eos: | ||
| case OgaFinishReason_StopSequence: |
There was a problem hiding this comment.
The exact non-release CI dependency (0.16.0-dev1001400138) defines OgaFinishReason_StopSequence, and the branch now compiles successfully against that package. Commit dd13bfe adds a preprocessor compatibility guard for both StopSequence and the later upstream StopString spelling, so either Engine-capable API variant maps to FOUNDRY_LOCAL_FINISH_STOP without referencing an undefined symbol.
Invalidate retained requests when explicit guidance changes, and close newly-created Engine conversations if turn setup or adapter construction fails. Accept both StopSequence and StopString finish-reason symbols while compiling against the exact non-release CI nightly. Files changed: chat_session.cc, onnx_engine_chat_generator.cc. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🔵 Needs a closer look
Engine special-token decoding is not model-agnostic, and static Engine behavior lacks runtime coverage.
Review details
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
sdk_v2/cpp/src/inferencing/generative/chat/onnx_engine_chat_generator.cc:99
- The Engine decoder only surfaces special tokens whose decoded spelling contains
tool_callorthink. The rest of this codebase treats BOT/EOT/BOR/EOR as model-configured token IDs (GetTagInfo()), so a valid model-specific marker with another spelling is silently decoded away and tool/reasoning parsing breaks on the Engine path. Match the configured tag IDs, as the Generator decoder already does, rather than guessing from token text.
sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc:707 - The static Engine branch has no runtime coverage: the new staging fixture always writes
dynamic_batching, and the only static test parses configuration. Static mode takes different paths here and in search-option setup, and rebuilds every turn, so add a staged static model test covering at least multi-turn reconstruction (ideally cancellation and concurrent sessions too).
- Files reviewed: 25/25 changed files
- Comments generated: 0 new
- Review effort level: Balanced
Decode model-configured BOT/EOT/BOR/EOR token IDs exactly like the Generator backend instead of matching English token spellings. Stage a static-batching variant of the shared model and verify that a rebuilt second turn retains committed conversation history. Files changed: onnx_engine_chat_generator.cc/.h, chat_session_test.cc. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
Dynamic capacity handling can hang requests, and the new Engine test staging is incompatible with the configured CI cache and model fixture.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 25/25 changed files
- Comments generated: 3
- Review effort level: Balanced
| if ((flags & OgaEngineEventFlag_Token) != 0) { | ||
| conversation->tokens.push_back(event->Token()); | ||
| ++conversation->sequence_length; | ||
| } | ||
| if ((flags & OgaEngineEventFlag_TurnFinished) != 0) { |
There was a problem hiding this comment.
Handled in 0d100b7. Request-less capacity events now evict a completed retained conversation and close its native request, freeing a dynamic batch slot. A later turn on that session receives a typed eviction signal and rebuilds from committed ChatSession history. Retryable events are also bounded so the worker cannot spin indefinitely.
| : path_(source.parent_path() / ("engine-chat-test-" + BackendName(backend) + "-" + | ||
| std::to_string(fl::test::CurrentPid()))) { | ||
| std::error_code ec; | ||
| std::filesystem::remove_all(path_, ec); | ||
| std::filesystem::create_directories(path_); |
There was a problem hiding this comment.
Handled in 0d100b7. EngineModelStaging now uses the existing TempPath helper under the writable system temp directory rather than writing beside FOUNDRY_TEST_DATA_DIR. Copied files are made owner-writable and the rewritten config stream is validated.
| if (backend == ChatBackendKind::kDynamicEngine) { | ||
| config["engine"] = { | ||
| {"dynamic_batching", {{"max_batch_size", 2}, {"max_scheduled_tokens", 2048}}}, | ||
| }; |
There was a problem hiding this comment.
Handled in 0d100b7. The generic CPU fixture is no longer mislabeled as dynamic: it is staged only with static_batching, which is compatible with its shared-buffer KV layout. The concurrency, cancellation, and multi-turn reconstruction tests now explicitly assert the static Engine backend. CI does not currently provision a paged-KV model, so this avoids claiming unsupported dynamic coverage.
Evict dormant dynamic conversations under capacity pressure so blocked turns can progress and rebuild evicted state from committed chat history. Stage the shared CPU model in a writable temp directory and exercise only its supported static Engine configuration. Files changed: chat_session.cc, onnx_chat_engine.cc, onnx_chat_engine.h, chat_session_test.cc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
Capacity eviction can race with initial turn startup and spuriously fail a newly created conversation.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 25/25 changed files
- Comments generated: 1
- Review effort level: Balanced
Only completed conversations with a nonzero turn ID are dormant candidates, preventing capacity pressure from closing a request between creation and its first BeginTurn. Files changed: onnx_chat_engine.cc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
The capacity-eviction and dynamic retained-conversation recovery path lacks automated coverage.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 25/25 changed files
- Comments generated: 1
- Review effort level: Balanced
| throw std::runtime_error("ORT GenAI Engine failed with error code " + | ||
| std::to_string(event->ErrorCode())); | ||
| } | ||
| if ((flags & OgaEngineEventFlag_CapacityBlocked) != 0 && EvictDormantConversation()) { |
There was a problem hiding this comment.
The requested dynamic integration test cannot be added to this CI suite with the models it provisions. .pipelines/templates/fetch-models-from-artifacts-feed.yml:12-22 supplies only qwen2.5-0.5b-instruct-generic-cpu-4 for this fixture, and its genai_config.json uses fixed past_key_values.%d.* / present.%d.* shared-buffer tensors rather than the paged-KV inputs required by dynamic batching. That is why 0d100b7 removed the invalid dynamic mutation and now explicitly asserts static Engine. Re-labeling this model dynamic would reproduce the prior invalid test, while adding a new model download violates the existing no-download fixture policy and requires a separately published paged-KV artifact. The capacity path is implemented against the documented Engine contract, but deterministic end-to-end coverage must follow when CI provisions a real paged-KV model. No safe test change is possible in this PR with the available fixture.
There was a problem hiding this comment.
🔵 Needs a closer look
The experimental native Engine lifecycle, concurrency, and unautomated dynamic-retention path warrant final human review.
Review details
- Files reviewed: 25/25 changed files
- Comments generated: 0 new
- Review effort level: Balanced
Summary
genai_config.jsonChatSessionsmax_batch_sizeEngine integration details
Missing assistant-turn boundary
ORT GenAI Engine does not append or retain a generated EOS token in the request's logical sequence. A later
BeginTurnmust therefore include the model-specific tokens that close the prior assistant turn before introducing the next user turn. The Generator path already retains EOS, so reusing its append fragment caused dynamic Engine sessions to lose conversational context even though the request and KV cache remained resident.This change derives the continuation boundary from the model's own chat template using a synthetic assistant marker, removes the marker, and submits the remaining boundary plus new-turn fragment. This avoids hard-coding Qwen token IDs or template syntax.
Per-turn output limits
The initial Engine turn received
max_output_tokens, but retained turns calledBeginTurnwithoutOgaTurnOptions. Foundry could stop reading at its external limit while the Engine still considered the turn active, making the next continuation invalid and allowing uncapped work when callers relied on the default limit.Every Engine turn now receives an explicit resolved limit, including the 2048-token text default. Foundry waits for the corresponding terminal event, consumes exact per-turn Engine usage, and recreates retained state when request-level sampling settings that are baked into
OgaGeneratorParamschange.Backend contract
engine.dynamic_batchingengine.static_batchingEngine selection is limited to text chat models. Multimodal Engine declarations fail at model load rather than failing later during a request.
Validation
mainEngine APIs4,6,8, and10sapphirewhile a peer request concurrently returned14Dependency assumption
This draft pins
Microsoft.ML.OnnxRuntimeGenAI.Foundryto 0.15.3. It assumes that release contains the validated Engine C/C++ API and the static batching capacity behavior currently available in the tested upstream build.