Enable model-selected Engine chat batching

sdk_v2/cpp/CMakeLists.txt

@@ -213,6 +213,8 @@ set(FOUNDRY_LOCAL_SOURCES
     src/inferencing/generative/audio/pcm_utils.cc
     src/inferencing/generative/embeddings/embeddings_session.cc
     src/inferencing/generative/chat/chat_generator.cc
+    src/inferencing/generative/chat/onnx_chat_engine.cc
+    src/inferencing/generative/chat/onnx_engine_chat_generator.cc
     src/inferencing/session/session.cc
     src/inferencing/session/session_manager.cc
     src/inferencing/generative/chat/chat_session.cc

sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.cc

@@ -4,6 +4,10 @@
 
 namespace fl {
 
+std::optional<ChatTurnUsage> ChatGenerator::GetTurnUsage() const {
+  return std::nullopt;
+}
+
 std::string ChatGenerator::GenerateAll() {
   std::string result;
 

sdk_v2/cpp/src/inferencing/generative/chat/chat_generator.h

@@ -3,9 +3,20 @@
 #pragma once
 
 #include <string>
+#include <optional>
+#include <vector>
 
 namespace fl {
 
+class GenAIModelInstance;
+struct MessageItem;
+struct SearchOptions;
+
+struct ChatTurnUsage {
+  int prompt_tokens = 0;
+  int generated_tokens = 0;
+};
+
 /// Abstract interface for token-by-token text generation.
 /// One generator per request — not reusable, not thread-safe.
 /// Follows the classic pull-based iterator pattern:
@@ -41,6 +52,21 @@ class ChatGenerator {
   /// After cancellation, IsDone() should return true on the next check.
   virtual void Cancel() = 0;
 
+  /// Append a new conversational turn to retained model state.
+  virtual int AppendMessages(const std::vector<MessageItem>& new_messages,
+                             GenAIModelInstance& model,
+                             const std::string& tools_json,
+                             const SearchOptions& options) = 0;
+
+  /// Returns whether this backend can rewind retained model state directly.
+  virtual bool CanRewind() const = 0;
+
+  /// Rewind retained model state to a prior token position.
+  virtual void RewindTo(int token_count) = 0;
+
+  /// Return exact usage for the most recently completed turn when the backend exposes it.
+  virtual std::optional<ChatTurnUsage> GetTurnUsage() const;
+
  protected:
   ChatGenerator() = default;
 };

sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc

@@ -5,6 +5,7 @@
 
 #include "contracts/chat_completions.h"
 #include "contracts/chat_completions_converter.h"
+#include "inferencing/generative/chat/onnx_engine_chat_generator.h"
 #include "inferencing/generative/chat/onnx_chat_generator.h"
 #include "inferencing/generative/chat/reasoning_stream_splitter.h"
 #include "inferencing/generative/genai_model_instance.h"
@@ -48,6 +49,17 @@ void ApplyToolChoiceToContext(std::optional<flToolChoice> tool_choice, ToolCallC
   }
 }
 
+std::unique_ptr<ChatGenerator> CreateTextChatGenerator(const std::vector<MessageItem>& messages,
+                                                       const SearchOptions& options,
+                                                       GenAIModelInstance& model,
+                                                       const ToolCallContext& tool_ctx) {
+  if (model.GetGenAIConfig().GetChatBackendKind() != ChatBackendKind::kGenerator) {
+    return OnnxEngineChatGenerator::Create(messages, options, model, tool_ctx);
+  }
+
+  return OnnxChatGenerator::Create(messages, options, model, tool_ctx, /*use_full_context=*/true);
+}
+
 }  // namespace
 
 ChatSession::ChatSession(const fl::Model& catalog_model, GenAIModelInstance& model, ILogger& logger, ITelemetry& telemetry)
@@ -71,7 +83,9 @@ ChatSession::ChatSession(ChatSession&& other) noexcept
       history_(std::move(other.history_)),
       turns_(std::move(other.turns_)),
       session_options_(std::move(other.session_options_)),
-      cached_generator_(std::move(other.cached_generator_)) {
+      cached_generator_(std::move(other.cached_generator_)),
+      cached_tool_ctx_(std::move(other.cached_tool_ctx_)),
+      cached_search_options_(std::move(other.cached_search_options_)) {
   other.owns_session_ = false;
 }
 
@@ -457,6 +471,12 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response)
   int prompt_tokens = 0;
   int pre_turn_token_count = 0;
 
+  if (cached_generator_ &&
+      !cached_search_options_.HasSameRetainedGenerationSettings(effective_options)) {
+    cached_generator_.reset();
+    cached_tool_ctx_ = {};
+  }
+
   if (cached_generator_) {
     // Check if guidance requirements changed since the generator was created. Guidance (LARK grammar) is baked into
     // the OGA generator at creation time and cannot be changed. If tool_choice went from "required" to "auto" (or
@@ -467,14 +487,17 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response)
     bool prev_needs_guidance = cached_tool_ctx_.tool_output && !cached_tool_ctx_.text_output;
     bool curr_needs_guidance = turn_tool_ctx.tool_output && !turn_tool_ctx.text_output;
 
-    if (prev_needs_guidance != curr_needs_guidance) {
+    const bool static_engine =
+        Model().GetGenAIConfig().GetChatBackendKind() == ChatBackendKind::kStaticEngine;
+    if (prev_needs_guidance != curr_needs_guidance || static_engine) {
       // Guidance requirements changed — invalidate. The branch below will rebuild from full history.
       cached_generator_.reset();
       cached_tool_ctx_ = {};
     } else {
       // Continuous decoding: append only the new messages to the existing generator.
       pre_turn_token_count = cached_generator_->TokenCount();
-      prompt_tokens = cached_generator_->AppendMessages(new_messages, Model(), cached_tool_ctx_.tools_json);
+      prompt_tokens =
+          cached_generator_->AppendMessages(new_messages, Model(), cached_tool_ctx_.tools_json, effective_options);
 
       // Refresh per-turn fields (tool_choice, guidance) while keeping session-level definitions stable.
       UpdateToolContextForTurn(request, cached_tool_ctx_);
@@ -491,7 +514,7 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response)
     all_messages.insert(all_messages.end(), history_.begin(), history_.end());
     all_messages.insert(all_messages.end(), new_messages.begin(), new_messages.end());
 
-    std::unique_ptr<OnnxChatGenerator> generator;
+    std::unique_ptr<ChatGenerator> generator;
     if (media_turn) {
       // Media is single-shot: the generator is dropped after the turn (see
       // CommitTurn cleanup below) because AppendMessages can't extend a
@@ -500,18 +523,18 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response)
       // gigabytes (262k tokens × 28 layers × 8 heads × 128 dims for
       // qwen3-vl-2b ≈ 120 GB). Bound it to prompt + max_output_tokens.
       generator = OnnxChatGenerator::CreateWithMedia(all_messages, effective_options, Model(), images, audios,
-                         tool_ctx, /*use_full_context*/ false);
+                                                     tool_ctx, /*use_full_context*/ false);
     } else {
-      generator = OnnxChatGenerator::Create(all_messages, effective_options, Model(), tool_ctx,
-                                            /*use_full_context*/ true);
+      generator = CreateTextChatGenerator(all_messages, effective_options, Model(), tool_ctx);
     }
     prompt_tokens = generator->PromptTokenCount();
 
     cached_generator_ = std::move(generator);
     cached_tool_ctx_ = std::move(tool_ctx);
+    cached_search_options_ = effective_options;
   }
 
-  int max_output = effective_options.max_output_tokens.value_or(0);
+  const int max_output = ResolveMaxOutputTokens(effective_options);
 
   // Generate token-by-token with optional streaming.
   // Check request.canceled each iteration — a streaming callback returning
@@ -608,17 +631,33 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response)
   emit_segments(splitter.Flush());
   flush_accumulator();
 
+  if (request.canceled) {
+    cached_generator_->Cancel();
+  }
+
   int total_tokens = cached_generator_->TokenCount();
+  if (const auto turn_usage = cached_generator_->GetTurnUsage()) {
+    prompt_tokens = turn_usage->prompt_tokens;
+    total_tokens = turn_usage->prompt_tokens + turn_usage->generated_tokens;
+  }
+  bool discard_generator = false;
 
   if (request.canceled) {
-    // Rewind the generator to undo this turn's input. The generator remains valid
-    // for the next attempt — the caller can re-send the same input.
-    cached_generator_->RewindTo(pre_turn_token_count);
+    if (cached_generator_->CanRewind()) {
+      cached_generator_->RewindTo(pre_turn_token_count);
+    } else {
+      discard_generator = true;
+    }
   }
 
   ProcessGeneratedOutput(std::move(text), cached_tool_ctx_, effective_options, request.canceled,
                          response, prompt_tokens, total_tokens, std::move(streamed_tool_calls));
 
+  if (discard_generator) {
+    cached_generator_.reset();
+    cached_tool_ctx_ = {};
+  }
+
   // Commit input messages + assistant reply to history only on success (not cancelled)
   if (!request.canceled) {
     // LARK grammar (tool-call-only mode) is a single-shot finite parse. If generation was truncated while grammar was
@@ -632,7 +671,9 @@ void ChatSession::ProcessRequestImpl(const Request& request, Response& response)
     bool grammar_was_active = cached_tool_ctx_.tool_output && !cached_tool_ctx_.text_output;
     bool reasoning_was_active = cached_tool_ctx_.supports_reasoning;
 
-    if (grammar_was_active || reasoning_was_active) {
+    const bool static_engine =
+        Model().GetGenAIConfig().GetChatBackendKind() == ChatBackendKind::kStaticEngine;
+    if (grammar_was_active || reasoning_was_active || static_engine) {
       cached_generator_.reset();
       cached_tool_ctx_ = {};
     }
@@ -713,7 +754,7 @@ void ChatSession::ProcessChatCompletionsJson(const std::string& request_json, co
   }
 
   // Create generator
-  auto generator = OnnxChatGenerator::Create(messages, options, Model(), tool_ctx);
+  auto generator = CreateTextChatGenerator(messages, options, Model(), tool_ctx);
   int prompt_tokens = generator->PromptTokenCount();
 
   auto streaming_callback = CreateCallbackHandler(original_request);
@@ -826,7 +867,15 @@ void ChatSession::ProcessChatCompletionsJson(const std::string& request_json, co
     emit_ready_calls(out.ready_calls);
   }
 
+  if (original_request.canceled) {
+    generator->Cancel();
+  }
+
   int total_tokens = generator->TokenCount();
+  if (const auto turn_usage = generator->GetTurnUsage()) {
+    prompt_tokens = turn_usage->prompt_tokens;
+    total_tokens = turn_usage->prompt_tokens + turn_usage->generated_tokens;
+  }
 
   // Process the generated output into response items (MessageItem, ToolCallItem, etc.)
   // This also updates finish_reason, and usage on the response. Streamed-parsed tool calls are reused so call_ids
@@ -909,8 +958,11 @@ void ChatSession::UndoTurns(size_t count) {
       // Undoing all turns — destroy the generator entirely
       cached_generator_.reset();
       cached_tool_ctx_ = {};
-    } else {
+    } else if (cached_generator_->CanRewind()) {
       cached_generator_->RewindTo(target.pre_turn_token_count);
+    } else {
+      cached_generator_.reset();
+      cached_tool_ctx_ = {};
     }
   }
 

sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h

@@ -16,7 +16,7 @@
 namespace fl {
 
 class GenAIModelInstance;
-class OnnxChatGenerator;
+class ChatGenerator;
 
 /// A chat session that maintains conversation history across turns.
 /// Designed for multi-turn conversations where message history accumulates
@@ -119,11 +119,14 @@ class ChatSession : public Session {
 
   // Cached generator for continuous decoding (non-JSON path only).
   // Null until first non-JSON ProcessRequestImpl call.
-  std::unique_ptr<OnnxChatGenerator> cached_generator_;
+  std::unique_ptr<ChatGenerator> cached_generator_;
 
   // Tool context used when creating the cached generator.
   // Reused for subsequent turns to maintain tool definition consistency.
   ToolCallContext cached_tool_ctx_;
+
+  // Search settings baked into the retained generator or Engine request.
+  SearchOptions cached_search_options_;
 };
 
 }  // namespace fl

sdk_v2/cpp/src/inferencing/generative/chat/chat_template.cc

@@ -63,6 +63,26 @@ std::string BuildChatPrompt(const std::vector<MessageItem>& messages,
   return model.GetPreprocessor().ApplyChatTemplate(messages_str.c_str(), tools_ptr, /*add_generation_prompt=*/true);
 }
 
+std::string BuildChatContinuationPrompt(const std::vector<MessageItem>& messages,
+                                        GenAIModelInstance& model,
+                                        const std::string& tools_json) {
+  constexpr std::string_view kAssistantMarker = "__foundry_engine_assistant_boundary__";
+
+  std::vector<MessageItem> marked_messages;
+  marked_messages.reserve(messages.size() + 1);
+  marked_messages.emplace_back(FOUNDRY_LOCAL_ROLE_ASSISTANT, std::string(kAssistantMarker));
+  marked_messages.insert(marked_messages.end(), messages.begin(), messages.end());
+
+  auto marked_prompt = BuildChatPrompt(marked_messages, model, tools_json);
+  const auto marker_position = marked_prompt.find(kAssistantMarker);
+  if (marker_position == std::string::npos) {
+    FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL,
+             "chat template did not preserve assistant content needed to build an Engine continuation");
+  }
+
+  return marked_prompt.substr(marker_position + kAssistantMarker.size());
+}
+
 std::unique_ptr<OgaSequences> EncodePrompt(const std::string& prompt,
                                            GenAIModelInstance& model) {
   return model.GetPreprocessor().Encode(prompt.c_str());

sdk_v2/cpp/src/inferencing/generative/chat/chat_template.h

@@ -39,6 +39,12 @@ std::string BuildChatPrompt(const std::vector<MessageItem>& messages,
                             GenAIModelInstance& model,
                             const std::string& tools_json = "");
 
+/// Build the fragment appended after an Engine-generated assistant response.
+/// Engine does not retain the generated EOS token, so this includes the template's assistant-turn boundary.
+std::string BuildChatContinuationPrompt(const std::vector<MessageItem>& messages,
+                                        GenAIModelInstance& model,
+                                        const std::string& tools_json = "");
+
 /// Encode a prompt string into token sequences using the model's shared tokenizer (thread-safe).
 /// Returns a unique_ptr to OgaSequences. Caller takes ownership.
 ///