[LLM] Extract text-serving contracts and helpers

Prepare a batching-backed serving path without changing current worker behavior.
Share the existing prompt-history planner and define experimental serving types
for subsequent runtime and transport changes.

Test Plan: All 8 existing worker prompt-planner cases passed. The serving
contract and planner headers passed standalone C++17 compilation. Behavioral
serving tests are introduced alongside their implementations in later commits.

Authored with AI assistance from Devmate.

Author: metascroy

CMakeLists.txt

@@ -1015,6 +1015,8 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM)
   list(APPEND _executorch_extensions extension_llm_batching
        extension_llm_batching_module
   )
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/serving)
+  list(APPEND _executorch_extensions extension_llm_serving)
 endif()
 
 if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL)

examples/llm_server/cpp/BUCK

@@ -13,6 +13,7 @@ fbcode_target(_kind = runtime.cxx_library,
     header_namespace = "executorch/examples/llm_server/cpp",
     exported_deps = [
         "//executorch/extension/llm/runner:runner_lib",
+        "//executorch/extension/llm/serving:prompt_history",
         "//pytorch/tokenizers:headers",
     ],
     external_deps = ["nlohmann_json"],

examples/llm_server/cpp/test_worker_prefill_plan.cpp

@@ -13,7 +13,6 @@
 #include <executorch/examples/llm_server/cpp/worker_prefill_plan.h>
 
 #include <cstdio>
-#include <cstring>
 #include <string>
 #include <vector>
 
@@ -30,15 +29,15 @@ void expect(
     size_t suffix_start,
     const char* reason) {
   bool ok = p.action == action && p.suffix_start == suffix_start &&
-      std::strcmp(p.reason, reason) == 0;
+      p.reason == reason;
   if (!ok) {
     ++g_failures;
     printf(
         "  [FAIL] %s: got action=%d suffix_start=%zu reason=%s\n",
         name,
         (int)p.action,
         p.suffix_start,
-        p.reason);
+        p.reason.c_str());
   } else {
     printf("  [PASS] %s\n", name);
   }

examples/llm_server/cpp/worker_prefill_plan.h

@@ -8,65 +8,4 @@
 
 #pragma once
 
-// Pure warm-resume prefill decision for the model worker (no model/session/ET
-// dependency, so it is unit-testable in isolation). Given a named session's
-// resident token ids (exactly the tokens currently in its KV/recurrent state),
-// the new request's prompt token ids, and whether the session is dirty, decide
-// whether to reset + full-prefill or to keep the state and prefill only the
-// suffix. The decision is exact-token (never string / retokenized text), so a
-// kSuffix plan is always a correct continuation; anything uncertain falls back
-// to kFull. See worker_loop.h for how the plan is executed.
-
-#include <cstddef>
-#include <cstdint>
-#include <vector>
-
-namespace executorch {
-namespace extension {
-namespace llm {
-
-struct PrefillPlan {
-  enum Action {
-    kFull, // reset + prefill the whole prompt
-    kSuffix // keep state, prefill prompt_ids[suffix_start:] at pos>0
-  } action;
-  size_t suffix_start; // index in prompt_ids where prefill begins (0 for kFull)
-  // Reported as session_reset_reason: "new" (no resident), "exact_prefix"
-  // (suffix reuse), "dirty", "mismatch", "equal" (prompt == resident).
-  const char* reason;
-};
-
-inline PrefillPlan plan_prefill(
-    const std::vector<uint64_t>& resident,
-    const std::vector<uint64_t>& prompt,
-    bool dirty) {
-  if (dirty) {
-    return {PrefillPlan::kFull, 0, "dirty"};
-  }
-  if (resident.empty()) {
-    return {PrefillPlan::kFull, 0, "new"};
-  }
-  if (prompt.size() < resident.size()) {
-    return {PrefillPlan::kFull, 0, "mismatch"};
-  }
-  for (size_t i = 0; i < resident.size(); ++i) {
-    if (prompt[i] != resident[i]) {
-      return {PrefillPlan::kFull, 0, "mismatch"};
-    }
-  }
-  if (prompt.size() == resident.size()) {
-    // Prompt is exactly the resident state (no new tokens). The ideal would be
-    // to skip prefill and decode straight from the session's pending token, but
-    // the LLMSession API exposes no "is there a valid pending token?" query, so
-    // we conservatively reset + full prefill rather than risk
-    // prefill_tokens([]) or decoding from a stale/absent pending. Rare in
-    // practice (a new turn adds tokens); a session pending-state query is a
-    // possible later optimization.
-    return {PrefillPlan::kFull, 0, "equal"};
-  }
-  return {PrefillPlan::kSuffix, resident.size(), "exact_prefix"};
-}
-
-} // namespace llm
-} // namespace extension
-} // namespace executorch
+#include <executorch/extension/llm/serving/prompt_history.h>

extension/llm/runner/targets.bzl

@@ -37,11 +37,13 @@ def define_common_targets():
             exported_headers = [
                 "model_metadata.h",
                 "stats.h",
+                "text_stream.h",
                 "util.h",
             ],
             visibility = ["PUBLIC"],
             exported_deps = [
                 ":constants",
+                "//pytorch/tokenizers:headers",
                  "//executorch/extension/module:module" + aten_suffix,
                  "//executorch/extension/tensor:tensor" + aten_suffix,
             ],
@@ -101,12 +103,23 @@ def define_common_targets():
         )
 
         runtime.cxx_library(
-            name = "multimodal_runner_lib" + aten_suffix,
+            name = "multimodal_input" + aten_suffix,
             exported_headers = [
                 "audio.h",
                 "image.h",
-                "wav_loader.h",
                 "multimodal_input.h",
+            ],
+            visibility = ["PUBLIC"],
+            exported_deps = [
+                "//executorch/extension/tensor:tensor" + aten_suffix,
+                "//executorch/runtime/platform:platform",
+            ],
+        )
+
+        runtime.cxx_library(
+            name = "multimodal_runner_lib" + aten_suffix,
+            exported_headers = [
+                "wav_loader.h",
                 "multimodal_runner.h",
                 "multimodal_prefiller.h",
                 "multimodal_decoder_runner.h",
@@ -118,6 +131,7 @@ def define_common_targets():
                 ":text_decoder_runner" + aten_suffix,
                 ":text_prefiller" + aten_suffix,
                 ":image_prefiller" + aten_suffix,
+                ":multimodal_input" + aten_suffix,
                 ":text_token_generator" + aten_suffix,
             ],
         )

extension/llm/serving/BUCK

@@ -0,0 +1,13 @@
+load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+# The legacy worker only needs the pure planner, not serving or batching types.
+fbcode_target(_kind = runtime.cxx_library,
+    name = "prompt_history",
+    exported_headers = ["prompt_history.h"],
+    exported_deps = ["//executorch/runtime/platform:platform"],
+    header_namespace = "executorch/extension/llm/serving",
+    visibility = ["PUBLIC"],
+)

extension/llm/serving/CMakeLists.txt

@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+add_library(extension_llm_serving INTERFACE)
+target_compile_features(extension_llm_serving INTERFACE cxx_std_17)
+target_include_directories(
+  extension_llm_serving INTERFACE ${_common_include_directories}
+)
+
+target_link_libraries(extension_llm_serving INTERFACE executorch_core)
+
+install(TARGETS extension_llm_serving EXPORT ExecuTorchTargets)
+install(FILES types.h prompt_history.h
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/serving
+)

extension/llm/serving/prompt_history.h

@@ -0,0 +1,67 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+// Pure exact-token continuation decision. `resident` is the history against
+// which the full prompt is compared. It can be committed KV/recurrent tokens,
+// or logical history including a pending token if the caller's continuation
+// path feeds that token before the suffix. Dirty or uncertain histories replay
+// the full prompt; this helper neither inspects nor changes execution state.
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <executorch/runtime/platform/compiler.h>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+
+struct ET_EXPERIMENTAL PrefillPlan {
+  enum Action {
+    kFull, // reset + prefill the whole prompt
+    kSuffix // keep state, prefill prompt_ids[suffix_start:] at pos>0
+  } action;
+  size_t suffix_start; // index in prompt_ids where prefill begins (0 for kFull)
+  // Reported as session_reset_reason: "new" (no resident), "exact_prefix"
+  // (suffix reuse), "dirty", "mismatch", "equal" (prompt == resident).
+  std::string reason;
+};
+
+ET_EXPERIMENTAL inline PrefillPlan plan_prefill(
+    const std::vector<uint64_t>& resident,
+    const std::vector<uint64_t>& prompt,
+    bool dirty) {
+  if (dirty) {
+    return {PrefillPlan::kFull, 0, "dirty"};
+  }
+  if (resident.empty()) {
+    return {PrefillPlan::kFull, 0, "new"};
+  }
+  if (prompt.size() < resident.size()) {
+    return {PrefillPlan::kFull, 0, "mismatch"};
+  }
+  for (size_t i = 0; i < resident.size(); ++i) {
+    if (prompt[i] != resident[i]) {
+      return {PrefillPlan::kFull, 0, "mismatch"};
+    }
+  }
+  if (prompt.size() == resident.size()) {
+    // Equal histories conservatively replay: this helper cannot establish that
+    // decoding without a non-empty suffix is supported or safe for the caller.
+    return {PrefillPlan::kFull, 0, "equal"};
+  }
+  return {PrefillPlan::kSuffix, resident.size(), "exact_prefix"};
+}
+
+} // namespace llm
+} // namespace extension
+} // namespace executorch

extension/llm/serving/types.h

@@ -0,0 +1,130 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <variant>
+#include <vector>
+
+#include <executorch/extension/llm/batching/types.h>
+#include <executorch/extension/llm/runner/multimodal_input.h>
+#include <executorch/runtime/platform/compiler.h>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+namespace serving {
+
+// The full prompt, not a session delta. Preparation encodes each text segment
+// separately and appends ID segments verbatim, in order, without adding
+// BOS/EOS. Segment boundaries therefore matter even between adjacent text
+// segments. Only text and token segments are currently supported; other
+// modalities, encoding failures, and an empty prepared prompt are rejected with
+// InvalidArgument before changing session history.
+struct ET_EXPERIMENTAL PromptInput {
+  std::vector<MultimodalInput> segments;
+};
+
+struct ET_EXPERIMENTAL GenerationOptions {
+  // Must be positive when set. Capped by the context remaining after the full
+  // prompt; a prompt leaving no generation room is invalid. Unset uses that
+  // remaining context, or the service default when the context limit is
+  // unknown.
+  std::optional<std::int32_t> max_new_tokens;
+  // Per-request policy: finite temperature >= 0, finite top_p in (0, 1],
+  // and top_k >= 0. Invalid options leave existing session history unchanged.
+  batching::SamplingParams sampling;
+  // Added to the service's default stop tokens. A matched token is retained in
+  // logical session history, but excluded from text and generated_token_ids.
+  std::vector<batching::Token> stop_tokens;
+  // Non-empty strings matched across decoded pieces. The match and everything
+  // after it are hidden; a match invalidates exact token replay and warm reuse.
+  std::vector<std::string> stop_strings;
+  // Forwarded to the executor for this generation, together with sampling.
+  std::optional<std::uint64_t> seed;
+};
+
+enum class ET_EXPERIMENTAL FinishReason {
+  Stop, // EOS, a stop token, or a stop string
+  Length, // Generation budget exhausted
+  Cancelled, // Request cancellation, session close/reset, or shutdown
+  Failed, // See TerminalEvent::error
+};
+
+enum class ET_EXPERIMENTAL ErrorCode {
+  InvalidArgument,
+  NotReady,
+  SessionNotFound,
+  SessionBusy,
+  CapacityExceeded,
+  Internal,
+};
+
+struct ET_EXPERIMENTAL ServingError {
+  ErrorCode code = ErrorCode::Internal;
+  std::string message;
+};
+
+struct ET_EXPERIMENTAL GenerationStats {
+  // Full prepared prompt size, including any reused prefix.
+  std::size_t prompt_tokens = 0;
+  // Tokens processed by text output, excluding EOS/stop tokens but including
+  // the token that completes a string stop. Later discarded tokens do not
+  // count. Neither this count nor visible text describes logical session
+  // history.
+  std::size_t completion_tokens = 0;
+  // Committed prompt prefix reused without execution. A pending prediction
+  // fed by this request counts as prefilled, not reused.
+  std::size_t reused_prompt_tokens = 0;
+  // Actually consumed prompt tokens; may be partial on cancellation or failure.
+  std::size_t prefilled_prompt_tokens = 0;
+  double prefill_ms = 0.0;
+  double decode_ms = 0.0;
+  double total_ms = 0.0;
+  std::string session_reset_reason;
+  // Original output IDs counted by completion_tokens, excluding EOS/stop
+  // tokens; not a retokenization of visible text or the full logical history
+  // suffix. Present only for Stop/Length completion without a matched string
+  // stop or rendering error. An empty vector is distinct from absent.
+  std::optional<std::vector<batching::Token>> generated_token_ids;
+};
+
+struct ET_EXPERIMENTAL TextEvent {
+  std::string text;
+};
+
+struct ET_EXPERIMENTAL TerminalEvent {
+  FinishReason finish_reason = FinishReason::Stop;
+  GenerationStats stats;
+  // Present exactly when finish_reason is Failed.
+  std::optional<ServingError> error;
+};
+
+// Each admitted request emits zero or more text events, then one terminal event
+// after final text flushing and session finalization. Synchronous rejection
+// emits no events. If the sink throws, it is not called again, even for the
+// terminal.
+using GenerationEvent ET_EXPERIMENTAL = std::variant<TextEvent, TerminalEvent>;
+
+struct ET_EXPERIMENTAL ServingInfo {
+  bool ready = false;
+  std::size_t max_context_length = 0;
+  std::size_t max_sessions = 0;
+  // Logical serving slots, including slots reserved during reset/reopen.
+  // This is not a count of physically live executor sessions.
+  std::size_t active_sessions = 0;
+};
+
+} // namespace serving
+} // namespace llm
+} // namespace extension
+} // namespace executorch